typescriptwebsocketjestjswsjest-websocket-mock

How to mock websocket when using ws library in typescript?


I am using typescript "typescript": "^5.7.3" for a nodejs backend project

and

"ws": "^8.18.0"
"jest": "^29.7.0",
"ts-jest": "^29.2.5",

I tried using the testing library

jest-websocket-mock

used mocks under __tests__ folder with

ws.ts

export { WebSocket as default } from 'mock-socket';

This didn' work because on function doesn't exist on mock-socket Websocket

TypeError: ws_2.on is not a function

Any other ideas on the best way to mock when using ws library ?

My unfinished test code looked like

MyClient.test.ts this would listen to some price api in real life

import WS from 'jest-websocket-mock';

describe('test', () => {
    let server: WS;

    beforeEach(async () => {
        server = new WS('ws://localhost:1234');
    });

    afterEach(() => {
        WS.clean();
    });

    it('example test ws', async () => {
        myClient = new MyClient({
            socketUrl: 'ws://localhost:1234',
        });
        myClient.listenForPrices();
    });
});

and somewhere else MyClient.ts

import WebSocket from 'ws';

export default class MyClient {
    // eslint-disable-next-line no-useless-constructor
    constructor(private readonly config: { solanaWebsocketUrl: string }) {}

    async listenForPrice() {
        try {
            const ws = new WebSocket(this.config.solanaWebsocketUrl);

            ws.on('open', () => {
                const subscriptionMessage = JSON.stringify({
                    jsonrpc: '2.0',
                    id: 1,
                    method: 'listenPrice',
                });
                ws.send(subscriptionMessage);
            });

            ws.on('message', message => {
                const data = JSON.parse(message.toString());
                console.log('newMessage', message);
            });

            ws.on('error', error => {
                console.log(error);
            });

            ws.on('close', async () => {
                await setTimeout(() => {
                    this.listenForPrice();
                }, 5000);
            });
        } catch (error) {
            await setTimeout(() => {
                this.listenForPrice();
            }, 5000);
        }
    }
}


Solution

  • There is no library solution for mocking ws Websocket for nodejs and supporting on function.

    I had to write a custom solution, a custom mock under __mocks__/ws.ts

    import { WebSocket as MockWebSocket } from 'mock-socket';
    
    class CustomMockWebSocket extends MockWebSocket {
        // Store event listeners
        listeners: { [key: string]: Function[] } = {};
        static sendMockFn = jest.fn();
    
        // eslint-disable-next-line no-useless-constructor
        constructor(url: string) {
            super(url);
        }
    
        // Override the `on()` method to support custom event listeners
        public on(event: string, callback: Function): void {
            // Initialize the event listener array if not already present
            if (!this.listeners[event]) {
                this.listeners[event] = [];
            }
            // Add the event listener
            this.listeners[event].push(callback);
    
            // Call the corresponding `super` method if needed
            if (event === 'open' && this.readyState === WebSocket.OPEN) {
                callback();
            }
        }
    
        // Override the dispatch of events so that the custom listeners are called
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        public trigger(event: string, ...args: any[]): void {
            // Call all listeners for this event
            if (this.listeners[event]) {
                this.listeners[event].forEach(listener => {
                    listener(...args);
                });
            }
        }
    
        // You can mock other WebSocket behavior if necessary (e.g., send, close)
        public send(data: string | Blob | ArrayBuffer | ArrayBufferView): void {
            console.log('Mock sending data:', data);
            CustomMockWebSocket.sendMockFn(...arguments);
        }
    
        // You can also mock the close event and trigger the listeners for 'close'
        public close(): void {
            console.log('Mock WebSocket closed');
            this.trigger('close');
        }
    }
    
    // Export the extended WebSocket class
    export { CustomMockWebSocket as default };
    
    

    With this solution I can test anything, can verify that my websocket client really send data, example:

            expect(CustomMockWebSocket.sendMockFn).toHaveBeenCalledTimes(1);
            expect(CustomMockWebSocket.sendMockFn.mock.calls[0]).toEqual([
                '{"jsonrpc":"2.0","id":1,"method":"blabla","params":[{"mentions":["test"]},{"commitment":"processed"}]}',
            ]);
    

    I can also use import WS from 'jest-websocket-mock'; to create a mock server like server = new WS(process.env.WSS_ENDPOINT as string);

    and send messages to my client mock websocket that I custom mocked above

    await server.connected;
    server.send('{"data": "blabla"}');