From 75edd5f42810c5160986d418b1d09752b596d577 Mon Sep 17 00:00:00 2001
From: Jungho Bang <jungho.bang@coinbase.com>
Date: Sat, 30 Sep 2023 01:34:22 +0900
Subject: [PATCH] fix tests

---
 .../src/connection/RxWebSocket.test.ts        | 72 +++++++++----------
 1 file changed, 36 insertions(+), 36 deletions(-)

diff --git a/packages/wallet-sdk/src/connection/RxWebSocket.test.ts b/packages/wallet-sdk/src/connection/RxWebSocket.test.ts
index b07498e828..b822933791 100644
--- a/packages/wallet-sdk/src/connection/RxWebSocket.test.ts
+++ b/packages/wallet-sdk/src/connection/RxWebSocket.test.ts
@@ -1,5 +1,4 @@
 import WS from 'jest-websocket-mock';
-import { Observable } from 'rxjs';
 
 import { ConnectionState, RxWebSocket } from './RxWebSocket';
 
@@ -17,27 +16,15 @@ describe('RxWebSocket', () => {
 
   describe('is connected', () => {
     test('@connect & @disconnect', async () => {
+      const connectionStateListener = jest.fn();
+      rxWS.setConnectionStateListener(connectionStateListener);
+
       const client = rxWS.connect();
 
-      expect(client).toBeInstanceOf(Observable);
-      await client.toPromise();
+      await client;
       await server.connected;
 
-      // @ts-expect-error test private methods
-      expect(rxWS.webSocket).toBeInstanceOf(WebSocket);
-      // @ts-expect-error test private methods
-      rxWS.connectionStateSubject
-        .subscribe({
-          next: (val) => {
-            //  Connected state
-            expect(val).toEqual(ConnectionState.CONNECTED);
-          },
-          closed: (val: ConnectionState) => {
-            //  Disconnected state
-            expect(val).toEqual(ConnectionState.DISCONNECTED);
-          },
-        })
-        .unsubscribe();
+      expect(connectionStateListener).toHaveBeenCalledWith(ConnectionState.CONNECTED);
 
       // Sends data
       const webSocketSendMock = jest
@@ -49,54 +36,67 @@ describe('RxWebSocket', () => {
 
       // Disconnects
       rxWS.disconnect();
+      expect(connectionStateListener).toHaveBeenCalledWith(ConnectionState.DISCONNECTED);
       // @ts-expect-error test private methods
       expect(rxWS.webSocket).toBe(null);
     });
 
-    test('@connectionState$ & @incomingData$', () => {
-      expect(rxWS.connectionState$).toBeInstanceOf(Observable);
-      expect(rxWS.incomingData$).toBeInstanceOf(Observable);
-    });
-
-    test('@incomingJSONData$', () => {
-      expect(rxWS.incomingJSONData$).toBeInstanceOf(Observable);
-    });
-
     describe('errors & event listeners', () => {
       afterEach(() => rxWS.disconnect());
 
       test('@connect throws error when connecting again', async () => {
         const client = rxWS.connect();
-        await client.toPromise();
+        await client;
 
-        await expect(rxWS.connect().toPromise()).rejects.toThrow('webSocket object is not null');
+        await expect(rxWS.connect()).rejects.toThrow('webSocket object is not null');
       });
 
       test('@connect throws error & fails to set websocket instance', async () => {
         const errorConnect = new RxWebSocket('');
 
-        await expect(errorConnect.connect().toPromise()).rejects.toThrow(
+        await expect(errorConnect.connect()).rejects.toThrow(
           "Failed to construct 'WebSocket': 1 argument required, but only 0 present."
         );
       });
 
       test('onclose event throws error', async () => {
         const client = rxWS.connect();
-        await client.toPromise();
+        await client;
         await server.connected;
         server.error();
 
-        await expect(rxWS.connect().toPromise()).rejects.toThrow('websocket error 1000: ');
+        await expect(rxWS.connect()).rejects.toThrow('websocket error 1000: ');
       });
 
       test('onmessage event emits message', async () => {
+        const incomingDataListener = jest.fn();
+        rxWS.setIncomingDataListener(incomingDataListener);
+
+        const client = rxWS.connect();
+        await client;
+        await server.connected;
+
+        const message = {
+          type: 'ServerMessageType',
+          data: 'hello world',
+        };
+
+        server.send(JSON.stringify(message));
+        expect(incomingDataListener).toHaveBeenCalledWith(message);
+      });
+
+      test('onmessage event emits heartbeat message', async () => {
+        const incomingDataListener = jest.fn();
+        rxWS.setIncomingDataListener(incomingDataListener);
+
         const client = rxWS.connect();
-        await client.toPromise();
+        await client;
         await server.connected;
 
-        // @ts-expect-error test private methods
-        rxWS.incomingDataSubject.subscribe((val) => expect(val).toEqual('hello world'));
-        server.send('hello world');
+        server.send('h');
+        expect(incomingDataListener).toHaveBeenCalledWith({
+          type: 'Heartbeat',
+        });
       });
     });
   });