vitest mock socket

A set of utilities and matchers to aid in mocking websocket servers in vitest.

Built on top of mock-socket and a refactored implementation of vitest-websocket-mock and jest-websocket-mock

Install

1npm install -D vitest-mock-socket

Usage

Import and instantiate the instance

1import { WebSocketServer } from 'vitest-mock-socket';
2
3const server = new WebSocketServer(url);

Connect a client to the same url

1const client = new WebSocket(url);

Wait for the server to connect

1await server.connected(); 

The server will record all messages it receives

1client.send('hello');

The server can also send messages to all connected clients

1server.send('hello everyone');

The server will also handle json messages out of the box

1server.send({ foo: 'bar' })

Simulate an error and close the connection

1server.error();

Gracefully close the connection

1server.close();

The instance also has a static method to gracefully close all open connections. This is particularly useful to reset the environment between test runs.

1WebSocketServer.clean();

WebSocketServer constructor

Methods

• connected
  • a Promise that resolves every time the mock server receives a new connection. The resolved value is the WebSocket client instance that initiated the connection.
• closed
  • a Promise that resolves every time a connection to the mock server is closed.
• nextMessage
  • a Promise that resolves every time a mock server instance receives a new message.
• send
  • send a message to all connected clients.
• close
  • gracefully closes all opened connections.
• error
  • sends an error message to all connected clients and closes all opened connections.
• on
  • attach event listeners to handle new connection, message and close events. The callback receives the socket as its only argument.

Options

The constructor accepts an optional options object as second argument.

The options supported by the mock-socket library are directly passed-through to the mock-server's constructor.

1export interface WebSocketServerOptions extends MockSocket.ServerOptions {}

1const server = new WebSocketServer(url, options);

verifyClient

A verifyClient function can be given in the options for the vitest-mock-socket constructor.

This can be used to test behavior for a client that connects to a WebSocket server that it is blacklisted from. For example:

Note : Currently mock-socket's implementation does not send any parameters to this function (unlike the real ws implementation).

1test('rejects connections that fail the verifyClient option', async () => {
2  new WebSocketServer('ws://localhost:1234', { verifyClient: () => false });
3  const errorCallback = vitest.fn();
4
5  await expect(
6    new Promise((resolve, reject) => {
7      errorCallback.mockImplementation(reject);
8      const client = new WebSocket(url);
9      client.onerror = errorCallback;
10      client.onopen = resolve;
11    })
12    // WebSocket onerror event gets called with an event of type error and not an error
13  ).rejects.toEqual(expect.objectContaining({ type: 'error' }));
14});

selectProtocol

A selectProtocol function can be given in the options for the vitest-mock-socket constructor. This can be used to test behaviour for a client that connects to a WebSocket server using the wrong protocol.

1test('rejects connections that fail the selectProtocol option', async () => {
2  const selectProtocol = () => null;
3  new WebSocketServer('ws://localhost:1234', { selectProtocol });
4  const errorCallback = vitest.fn();
5
6  await expect(
7    new Promise((resolve, reject) => {
8      errorCallback.mockImplementationOnce(reject);
9      const client = new WebSocket('ws://localhost:1234', 'foo');
10      client.onerror = errorCallback;
11      client.onopen = resolve;
12    })
13  ).rejects.toEqual(
14    // WebSocket onerror event gets called with an event of type error and not an error
15    expect.objectContaining({
16      type: 'error',
17      currentTarget: expect.objectContaining({ protocol: 'foo' }),
18    })
19  );
20});

Vitest matchers

Custom vitest matchers are included to ease running assertions on received messages:

.toReceiveMessage

An async matcher that waits for the next message received by the the mocked websocket server, and asserts its content. It will time out with a helpful message after 1000ms.

1test('the server keeps track of received messages, and yields them as they come in', async () => {
2  const server = new WebSocketServer(url);
3  const client = new WebSocket(url);
4
5  await server.connected();
6
7  client.send('hello');
8
9  await expect(server).toReceiveMessage('hello');  
10});

.toHaveReceivedMessages

A synchronous matcher that checks that all the expected messages have been received by the mock websocket server.

Note: Since this matcher is synchronous, there are situations where you must call await server.nextMessage() for each message sent before asserting with this matcher. You can get the same behavior without the need to manually call .nextMessage by using the asynchronous variant toHaveResolvedMessages

1test('the server keeps track of received messages, and yields them as they come in', async () => {
2  const server = new WebSocketServer(url);
3  const client = new WebSocket(url);
4
5  await server.connected();
6
7  client.send('hello');
8  client.send('goodbye');
9
10  await server.nextMessage();
11  await server.nextMessage();
12
13  expect(server).toHaveReceivedMessages(['hello', 'goodbye']);
14});

1test('server handles mixed message types', async () => {
2  const server = new WebSocketServer(url);
3  const client = new WebSocket(url);
4  
5  await server.connected();
6
7  client.send('hello there');
8  client.send(`{"type":"GREETING","payload":"how are you?"}`);
9  client.send(`{"type":"GREETING","payload":"good?"}`);
10
11  await server.nextMessage();
12  await server.nextMessage();
13  await server.nextMessage();
14
15  expect(server).toHaveReceivedMessages([
16    'hello there',
17    { type: 'GREETING', payload: 'how are you?' },
18    { type: 'GREETING', payload: 'good?' },
19  ]);
20});

.toHaveResolvedMessages

A asynchronous version of toHaveReceivedMessages. It will automatically resolve the message queue so you do not have to manually call server.nextMessage.

Default behavior is to match exactly.

1test('the server keeps track of received messages, and yields them as they come in', async () => {
2  const server = new WebSocketServer(url);
3  const client = new WebSocket(url);
4
5  await server.connected();
6
7  client.send('hello');
8  client.send('goodbye');
9
10  await expect(server).toHaveResolvedMessages(['hello', 'goodbye']);
11});

This would fail

1test('the server keeps track of received messages, and yields them as they come in.', async () => {
2  const server = new WebSocketServer(url);
3  const client = new WebSocket(url);
4
5  await server.connected();
6
7  client.send('hello');
8  client.send('goodbye');
9
10  await expect(server).toHaveResolvedMessages(['hello' ]);
11});

Accepts an options object to allow for partial matches.

1test('the server keeps track of received messages, and yields them as they come in.', async () => {
2  const server = new WebSocketServer(url);
3  const client = new WebSocket(url);
4
5  await server.connected();
6
7  client.send('hello');
8  client.send('goodbye');
9
10  await expect(server).toHaveResolvedMessages(['hello'], { partial: true });
11});

1  test('server handles mixed message types', async () => {
2    const server = new WebSocketServer(url);
3    const client = new WebSocket(url);
4    await server.connected();
5    client.send('hello there');
6    client.send(`{"type":"GREETING","payload":"how are you?"}`);
7    client.send(`{"type":"GREETING","payload":"good?"}`);
8    await expect(server).toHaveResolvedMessages([
9      'hello there',
10      { type: 'GREETING', payload: 'how are you?' },
11      { type: 'GREETING', payload: 'good?' },
12    ]);
13  });

Other examples

Send messages to multiple connected clients

1test('the mock server sends messages to connected clients', async () => {
2  const server = new WebSocketServer(url);
3
4  const client1 = new WebSocket(url);
5  await server.connected();
6  
7  const client2 = new WebSocket(url);
8  await server.connected();
9
10  const messages = { client1: [], client2: [] };
11  client1.onmessage = (e) => {
12    messages.client1.push(e.data);
13  };
14  client2.onmessage = (e) => {
15    messages.client2.push(e.data);
16  };
17
18  server.send('hello everyone');
19  expect(messages).toEqual({
20    client1: ['hello everyone'],
21    client2: ['hello everyone'],
22  });
23});

Sending errors

1test('the mock server sends errors to connected clients', async () => {
2  const server = new WebSocketServer(url);
3  const client = new WebSocket(url);
4  await server.connected();
5
6  let disconnected = false;
7  let error = null;
8  client.onclose = () => {
9    disconnected = true;
10  };
11  client.onerror = (e) => {
12    error = e;
13  };
14
15  server.send('hello everyone');
16  server.error();
17  expect(disconnected).toBe(true);
18  expect(error.origin).toBe('ws://localhost:1234/');
19  expect(error.type).toBe('error');
20});

Refuse connections example:

1it('the server can refuse connections', async () => {
2  const server = new WebSocketServer(url);
3  server.on('connection', (socket) => {
4    socket.close({ wasClean: false, code: 1003, reason: 'NOPE' });
5  });
6
7  const client = new WebSocket(url);
8  client.onclose = (event: CloseEvent) => {
9    expect(event.code).toBe(1003);
10    expect(event.wasClean).toBe(false);
11    expect(event.reason).toBe('NOPE');
12  };
13
14  expect(client.readyState).toBe(WebSocket.CONNECTING);
15
16  await server.connected();
17  expect(client.readyState).toBe(WebSocket.CLOSING);
18
19  await server.closed();
20  expect(client.readyState).toBe(WebSocket.CLOSED);
21});

Environment set up and tear down between tests

You can set up a mock server and a client, and reset them between tests:

1beforeEach(async () => {
2  server = new WebSocketServer(url);
3  client = new WebSocket(url);
4  await server.connected();
5});
6
7afterEach(() => {
8  WebSocketServer.clean();
9});

Known issues

mock-socket has a strong usage of delays (setTimeout to be more specific). This means using vi.useFakeTimers(); will cause issues such as the client appearing to never connect to the server.

While running the websocket server from tests within the vitest-dom environment (as opposed to node) you may see errors of the nature:

1 ReferenceError: setImmediate is not defined

You can work around this by installing the setImmediate shim from https://github.com/YuzuJS/setImmediate and adding require('setimmediate'); to your setupTests.js.

Using vitest-mock-socket to interact with a non-global WebSocket object

vitest-mock-socket uses the mock-socket library. under the hood to mock out WebSocket clients.

Out of the box, mock-socket will only mock out the global WebSocket object. If you are using a third-party WebSocket client library (eg. a Node.js implementation, like ws), you'll need to set up a manual mock:

• Create a __mocks__ folder in your project root
• Add a new file in the __mocks__ folder named after the library you want to mock out. For instance, for the ws library: __mocks__/ws.js.
• Export Mock Socket's implementation in-lieu of the normal export from the library you want to mock out. For instance, for the ws library:

1// __mocks__/ws.js
2
3export { WebSocket as default } from 'mock-socket';

• Somewhere in the test files, call vi.mock with the name of the library you want to mock. For instance, for the ws library:

1// example.test.js
2import WebSocket from 'ws';
3import { vi } from 'vitest';
4
5vi.mock('ws');
6
7// do some tests...

NOTE The ws library is not 100% compatible with the browser API, and vitest-mock-socket's dependency mock-socket only implements the browser API.

As a result, vitest-mock-socket will only work with the ws library if you restrict yourself to the browser APIs.