I have this example code

Provider.ts

import WebSocket from 'ws';

class Provider {
    async listen(onNewData: (data: any) => void) {
        ws.on('open', () => {
            const subscriptionMessage = JSON.stringify({
                jsonrpc: '2.0',
                id: 1,
                method: 'pricesSubscribe',
            });
            ws.send(subscriptionMessage);
        });

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

I want to write a test that

1. asserts that .on('open') is called with

JSON.stringify({
        jsonrpc: '2.0',
        id: 1,
        method: 'pricesSubscribe',
    })

2. I send the message myself which is received into on('message')
3. I want to assert that newData callback is called with the proper argument

I have this example code

Provider.ts

import WebSocket from 'ws';

class Provider {
    async listen(onNewData: (data: any) => void) {
        ws.on('open', () => {
            const subscriptionMessage = JSON.stringify({
                jsonrpc: '2.0',
                id: 1,
                method: 'pricesSubscribe',
            });
            ws.send(subscriptionMessage);
        });

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

I want to write a test that

1. asserts that .on('open') is called with

JSON.stringify({
        jsonrpc: '2.0',
        id: 1,
        method: 'pricesSubscribe',
    })

2. I send the message myself which is received into on('message')
3. I want to assert that newData callback is called with the proper argument

• typescript
• websocket
• jestjs
• mocking
• ws

2 Answers 2

Using ws to communicate via a ws WebSocket in code which can be embedded in a unit test:

import WebSocket, { WebSocketServer } from 'ws';

class ChatApp {
    constructor(url) {
        this.messages = [];
        this.connection = new WebSocket(url);

        this.connection.on('message', (data) => {
            this.messages.push(data.toString());
        });
    }

    sendMessage(message) {
        this.connection.send(message);
    }

    static connect(url) {
        const chatApp = new ChatApp(url);
        return new Promise((resolve, reject) => {
            chatApp.connection.on('open', () => resolve(chatApp));
            chatApp.connection.on('error', (err) => reject(err));
        });
    }
}

function sleep(delay) {
    return new Promise((resolve) => setTimeout(resolve, delay));
}

async function run() {
    console.log('Init server...');
    const fakeURL = 'ws://localhost:8080';

    const wss = new WebSocketServer({ port: 8080 });
    wss.on('connection', (socket) => {
        socket.on('message', (message) => {
            console.log(`received message from socket`);
            const { jsonrpc, id, method } = JSON.parse(message.toString());
            console.log(`Received: jsonrpc=${jsonrpc}, id=${id}, method=${method}`);
        });
    });

    const app = await ChatApp.connect(fakeURL);
    const subscriptionMessage = JSON.stringify({
        jsonrpc: '2.0',
        id: 1,
        method: 'pricesSubscribe',
    });
    app.sendMessage(subscriptionMessage);

    await sleep(100);
    console.log(`Stop server`);
    wss.close();
}

run();

Outputs (Node.js, ESM module):

Init server...
received message from socket
Received: jsonrpc=2.0, id=1, method=pricesSubscribe
Stop server

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

I also used jest-websocket-mock only for the server The WebSocket from mock-server doesn't contain on function and fails on backend nodejs tests, that's why I created this custom mock

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 };