From 126780840b6b79d277607c15f0cfbfbb0b3dc08f Mon Sep 17 00:00:00 2001
From: Badi Ifaoui <badi1wow@gmail.com>
Date: Thu, 13 Jul 2023 10:09:12 +0100
Subject: [PATCH] chore: websocket client unit tests

---
 src/lib/WebSocketClient.ts          |  15 ++++
 tests/unit/websocket-client.spec.js | 106 ++++++++++++++++++++++++++++
 2 files changed, 121 insertions(+)
 create mode 100644 tests/unit/websocket-client.spec.js

diff --git a/src/lib/WebSocketClient.ts b/src/lib/WebSocketClient.ts
index 8f1969a6b..1126578ff 100644
--- a/src/lib/WebSocketClient.ts
+++ b/src/lib/WebSocketClient.ts
@@ -76,6 +76,21 @@ class WebSocketClient {
     this.subscribers[message.payload][uuid] = callback;
     return () => {
       delete this.subscribers[message.payload][uuid];
+      if (Object.keys(this.subscribers[message.payload]).length === 0) {
+        // should remove the message from the queue if there are no subscribers
+        this.subscribersQueue = this.subscribersQueue.filter(
+          (msg) => msg.payload !== message.payload,
+        );
+
+        // should unsubscribe from the channel if there are no subscribers
+        Object.keys(WEB_SOCKET_SOURCE).forEach((source) => {
+          this.wsClient.send(JSON.stringify({
+            ...message,
+            op: WEB_SOCKET_UNSUBSCRIBE,
+            source,
+          }));
+        });
+      }
     };
   }
 
diff --git a/tests/unit/websocket-client.spec.js b/tests/unit/websocket-client.spec.js
new file mode 100644
index 000000000..c5c9ec54d
--- /dev/null
+++ b/tests/unit/websocket-client.spec.js
@@ -0,0 +1,106 @@
+import WebSocketClient from '../../src/lib/WebSocketClient';
+import {
+  NETWORK_TESTNET,
+  WEB_SOCKET_SOURCE,
+  WEB_SOCKET_CHANNELS,
+} from '../../src/popup/utils';
+
+const testAddress = 'ak_2fxchiLvnj9VADMAXHBiKPsaCEsTFehAspcmWJ3ZzF3pFK1hB5';
+
+const WEB_SOCKET_SOURCE_COUNT = Object.keys(WEB_SOCKET_SOURCE).length;
+
+describe('WebSocketClient', () => {
+  beforeAll(async () => {
+    WebSocketClient.connect(NETWORK_TESTNET.websocketUrl);
+    await new Promise((resolve) => {
+      WebSocketClient.wsClient.onopen = () => {
+        resolve();
+      };
+    });
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  describe('Can connect to ', () => {
+    it('should connect and set isWsConnected to true', () => {
+      expect(WebSocketClient.isWsConnected).toBe(true);
+    });
+  });
+
+  describe('Can handle transactions subscription', () => {
+    let transactionChannelUnsubscribe;
+    it('should subscribe for transactions', () => {
+      const mockSend = jest.spyOn(WebSocketClient.wsClient, 'send');
+      transactionChannelUnsubscribe = WebSocketClient.subscribeForTransactionsUpdates(
+        jest.fn(),
+      );
+      expect(mockSend).toHaveBeenCalledTimes(WEB_SOCKET_SOURCE_COUNT);
+      expect(Object.keys(WebSocketClient.subscribers[WEB_SOCKET_CHANNELS.Transactions]))
+        .toHaveLength(1);
+      expect(WebSocketClient.subscribersQueue)
+        .toHaveLength(1);
+    });
+
+    it('should unsubscribe from transactions', () => {
+      const mockSend = jest.spyOn(WebSocketClient.wsClient, 'send');
+      transactionChannelUnsubscribe();
+      expect(mockSend).toHaveBeenCalledTimes(WEB_SOCKET_SOURCE_COUNT);
+      expect(Object.keys(WebSocketClient.subscribers[WEB_SOCKET_CHANNELS.Transactions]))
+        .toHaveLength(0);
+      expect(WebSocketClient.subscribersQueue)
+        .toHaveLength(0);
+    });
+  });
+
+  describe('Can handle micro blocks subscription', () => {
+    let microBlocksChannelUnsubscribe;
+    it('should subscribe for micro blocks', () => {
+      const mockSend = jest.spyOn(WebSocketClient.wsClient, 'send');
+      microBlocksChannelUnsubscribe = WebSocketClient.subscribeForMicroBlocksUpdates(
+        jest.fn(),
+      );
+      expect(mockSend).toHaveBeenCalledTimes(WEB_SOCKET_SOURCE_COUNT);
+      expect(Object.keys(WebSocketClient.subscribers[WEB_SOCKET_CHANNELS.MicroBlocks]))
+        .toHaveLength(1);
+      expect(WebSocketClient.subscribersQueue)
+        .toHaveLength(1);
+    });
+
+    it('should unsubscribe from micro blocks', () => {
+      const mockSend = jest.spyOn(WebSocketClient.wsClient, 'send');
+      microBlocksChannelUnsubscribe();
+      expect(mockSend).toHaveBeenCalledTimes(WEB_SOCKET_SOURCE_COUNT);
+      expect(Object.keys(WebSocketClient.subscribers[WEB_SOCKET_CHANNELS.MicroBlocks]))
+        .toHaveLength(0);
+      expect(WebSocketClient.subscribersQueue)
+        .toHaveLength(0);
+    });
+  });
+
+  describe('Can handle account transactions subscription', () => {
+    let transactionChannelUnsubscribe;
+    it('should subscribe for account transactions', () => {
+      const mockSend = jest.spyOn(WebSocketClient.wsClient, 'send');
+      transactionChannelUnsubscribe = WebSocketClient.subscribeForAccountUpdates(
+        testAddress, jest.fn(),
+      );
+      expect(mockSend).toHaveBeenCalledTimes(WEB_SOCKET_SOURCE_COUNT);
+      expect(Object.keys(WebSocketClient.subscribers[WEB_SOCKET_CHANNELS.Object]))
+        .toHaveLength(1);
+      expect(WebSocketClient.subscribersQueue)
+        .toHaveLength(1);
+    });
+
+    it('should unsubscribe from account transactions', () => {
+      const mockSend = jest.spyOn(WebSocketClient.wsClient, 'send');
+      transactionChannelUnsubscribe();
+      expect(mockSend).toHaveBeenCalledTimes(WEB_SOCKET_SOURCE_COUNT);
+      expect(Object.keys(WebSocketClient.subscribers[WEB_SOCKET_CHANNELS.Object]))
+        .toHaveLength(0);
+      expect(WebSocketClient.subscribersQueue)
+        .toHaveLength(0);
+    });
+  });
+});