使用 jest 的 websocket 服务器“ws”库的单元测试用例

Question

I am trying to figure out how to write unit test cases for Websocket server which is using the ws library. I did go through jest-websocket-mock but I think this is for browser based APIs and I want to test server using JEST .

Basic Code: Server.js

import { createServer } from 'https';
import { WebSocketServer } from 'ws';
import { readFileSync } from 'fs';

const server = createServer({
  cert: readFileSync(config.certs.sslCertPath),
  key: readFileSync(config.certs.sslKeyPath),
});

const wss = new WebSocketServer({ noServer: true });

server.on('upgrade', (request, socket, head) => {
  const origin = request && request.headers && request.headers.origin;
  const corsRegex = <regex>;

  if (origin.match(corsRegex) != null) {
    wss.handleUpgrade(request, socket, head, (ws) => {
      wss.emit('connection', ws, request);
    });
  } else {
    socket.destroy();
  }
});

wss.on('connection', (ws, req) => {
  ws.on('message', (messageg) => {
    try {
      console.log(message);
    } catch (error) {
      console.log(error);
    }
  });

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

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

Can someone please help me with how can I test the original server?

Answer 1

you need to create some kind of dependency injection mechanism here lets for example move all the socket initialization logic into a separate function

function initSocketEvents(wss) {
    wss.on('connection', (ws, req) => {
        ws.on('message', (messageg) => {
            try {
                console.log(message);
            } catch (error) {
                console.log(error);
            }
        });

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

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

now at the server initilization just call the function from a different file

...

const {initSocketEvents}  = require("./socket-handler") 
const wss = new WebSocketServer({ noServer: true });
initSocketEvents(was);

...

everything stays the same except the fact its way much easier to test it now at the test file

const {initSocketEvents}  = require("./socket-handler") 
const { assert } = require('console');
const { EventEmitter } = require('events');
class MyTestWebSocket extends EventEmitter { }
const mockWSS = new MyTestWebSocket()
initSocketEvents(mockWSS)
mockWSS.emit('connection', mockWSS)
assert(mockWSS.listenerCount('connection')===1)
assert(mockWSS.listenerCount('message')===1)
assert(mockWSS.listenerCount('close')===1)
assert(mockWSS.listenerCount('error')===1)

now it should be straightforward to separate each listener's logic and inject it outside the function and then assert the desired logic.