如何为HTTP Server升级与WebSocket连接编写单元测试

核心思路 #

要覆盖upgradeHandler和connectionHandler内部逻辑，关键是手动触发HTTP Server的upgrade事件，同时Mock WebSocket相关依赖，避免真实网络连接。

具体实现步骤 #

1. Mock HTTP Server与WebSocket模块 #

用Jest的jest.mock替换原生http和ws模块，或者Spy相关方法：

// 引入依赖
const http = require('http');
const WebSocket = require('ws');
const ClassB = require('./your-class-b-path');

// Mock HTTP Server
jest.mock('http', () => ({
  createServer: jest.fn(() => ({
    on: jest.fn(),
    listen: jest.fn((port, cb) => cb()),
    close: jest.fn((cb) => cb())
  }))
}));

// Mock WebSocket模块
jest.mock('ws', () => {
  const mockWsClient = jest.fn().mockImplementation(() => ({
    on: jest.fn(),
    send: jest.fn(),
    close: jest.fn()
  }));
  mockWsClient.Server = jest.fn().mockImplementation(() => ({
    on: jest.fn(),
    emit: jest.fn()
  }));
  return mockWsClient;
});

2. 手动触发upgrade事件 #

在测试用例中，提取HTTP Server的upgrade事件监听函数并手动调用：

describe('ClassB WebSocket Logic', () => {
  let classBInstance;
  let mockHttpServer;
  let upgradeHandler;

  beforeEach(() => {
    jest.clearAllMocks();
    classBInstance = new ClassB();
    mockHttpServer = http.createServer.mock.results[0].value;
    
    // 提取upgrade事件处理函数
    const onCalls = mockHttpServer.on.mock.calls;
    upgradeHandler = onCalls.find(call => call[0] === 'upgrade')[1];
  });

  test('should process upgrade event correctly', () => {
    // 模拟请求、socket、head参数
    const mockReq = { url: '/ws' };
    const mockSocket = {
      on: jest.fn(),
      write: jest.fn(),
      destroy: jest.fn()
    };
    const mockHead = Buffer.from('');

    // 手动触发upgrade事件
    upgradeHandler(mockReq, mockSocket, mockHead);

    // 验证WebSocket服务是否初始化
    expect(WebSocket.Server).toHaveBeenCalled();
    // 验证connection事件是否绑定处理函数
    const mockWsServer = WebSocket.Server.mock.results[0].value;
    expect(mockWsServer.on).toHaveBeenCalledWith('connection', expect.any(Function));
  });
});

3. 覆盖connectionHandler内部逻辑 #

提取WebSocket Server的connection事件处理函数，手动触发并验证业务逻辑：

test('should handle WebSocket connection properly', () => {
  // 先触发upgrade初始化WebSocket服务
  const mockReq = { url: '/ws' };
  const mockSocket = { on: jest.fn(), write: jest.fn(), destroy: jest.fn() };
  const mockHead = Buffer.from('');
  upgradeHandler(mockReq, mockSocket, mockHead);

  // 获取connection处理函数
  const mockWsServer = WebSocket.Server.mock.results[0].value;
  const connectionHandler = mockWsServer.on.mock.calls.find(call => call[0] === 'connection')[1];

  // 模拟客户端连接
  const mockClientWs = new WebSocket();
  connectionHandler(mockClientWs);

  // 验证message事件监听是否绑定
  expect(mockClientWs.on).toHaveBeenCalledWith('message', expect.any(Function));
  
  // 模拟message事件触发，验证处理逻辑
  const messageHandler = mockClientWs.on.mock.calls.find(call => call[0] === 'message')[1];
  const testMsg = Buffer.from('test payload');
  messageHandler(testMsg);

  // 根据业务逻辑添加断言，比如是否回复消息
  expect(mockClientWs.send).toHaveBeenCalledWith(expect.any(String));
});

4. 测试错误场景 #

模拟异常情况，验证错误处理逻辑：

test('should reject invalid upgrade requests', () => {
  const mockReq = { url: '/invalid-ws-path' };
  const mockSocket = { write: jest.fn(), destroy: jest.fn() };
  const mockHead = Buffer.from('');

  upgradeHandler(mockReq, mockSocket, mockHead);

  // 验证是否返回错误并关闭连接
  expect(mockSocket.write).toHaveBeenCalledWith('HTTP/1.1 400 Bad Request\r\n\r\n');
  expect(mockSocket.destroy).toHaveBeenCalled();
});

关键注意点 #

• 每次测试前调用jest.clearAllMocks()，避免用例间状态污染
• 若抽象类A的方法需要验证调用，可使用jest.spyOn(ClassB.prototype, 'abstractMethod')监听
• 若WebSocket使用自定义子协议，需在mockReq.headers中添加对应字段模拟请求头

内容的提问来源于stack exchange，提问作者preadeep xsy