test(server): 添加 server 模块单元测试

- 新增 vitest 配置和测试基础设施 - 添加 adapter.test.ts: 测试 Core 模块初始化和 Agent 管理 (18 tests) - 添加 token.test.ts: 测试 Token 生成、验证和中间件 (25 tests) - 添加 handler.test.ts: 测试权限处理器 (18 tests) - 添加 ws.test.ts: 测试 WebSocket 连接和消息处理 (19 tests) - 添加 sse.test.ts: 测试 SSE 事件发送 (14 tests) - 添加 sessions.test.ts: 测试会话路由 (16 tests) - 添加 config.test.ts: 测试配置路由 (10 tests) - 添加 context.test.ts: 测试上下文压缩路由 (9 tests) - 添加 providers.test.ts: 测试 Provider 管理路由 (18 tests) - 添加 manager.test.ts: 测试 SessionManager (48 tests) 总计 195 个测试，覆盖率从 0% 提升至 29.59%
2025-12-15 00:07:32 +08:00
parent 503e4c4ccd
commit 5835799b69
17 changed files with 3258 additions and 3 deletions
@@ -0,0 +1,346 @@
+/**
+ * WebSocket Handler 测试
+ *
+ * 测试 WebSocket 连接处理、消息路由、广播功能等
+ */
+
+import { describe, it, expect, beforeEach, vi } from 'vitest';
+
+// Create mock functions
+const mockExists = vi.fn();
+const mockUpdateStatus = vi.fn();
+const mockAddMessage = vi.fn();
+const mockGet = vi.fn();
+
+// Mock dependencies before imports
+vi.mock('../../src/session/manager.js', () => ({
+  getSessionManager: vi.fn(() => ({
+    exists: mockExists,
+    updateStatus: mockUpdateStatus,
+    addMessage: mockAddMessage,
+    get: mockGet,
+  })),
+}));
+
+vi.mock('../../src/agent/index.js', () => ({
+  processMessage: vi.fn().mockResolvedValue(undefined),
+  cancelProcessing: vi.fn(),
+}));
+
+vi.mock('../../src/permission/handler.js', () => ({
+  handlePermissionResponse: vi.fn().mockReturnValue(true),
+}));
+
+import {
+  handleWebSocket,
+  handleWebSocketMessage,
+  handleWebSocketClose,
+  broadcastToSession,
+  getSessionConnections,
+  getConnectionStats,
+} from '../../src/ws.js';
+import { processMessage, cancelProcessing } from '../../src/agent/index.js';
+import { handlePermissionResponse } from '../../src/permission/handler.js';
+
+// Mock WSContext
+function createMockWSContext() {
+  return {
+    send: vi.fn(),
+    close: vi.fn(),
+  };
+}
+
+// Counter to generate unique session IDs for test isolation
+let testCounter = 0;
+function getUniqueSessionId(prefix = 'session') {
+  return `${prefix}-${Date.now()}-${testCounter++}`;
+}
+
+describe('WebSocket Handler', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockExists.mockReturnValue(false);
+    mockAddMessage.mockReturnValue({ id: 'msg-1', role: 'user', content: '', timestamp: Date.now() });
+  });
+
+  describe('handleWebSocket - 连接处理', () => {
+    it('无效 session 发送错误并关闭连接', () => {
+      const ws = createMockWSContext();
+      mockExists.mockReturnValue(false);
+
+      handleWebSocket(ws as any, 'invalid-session');
+
+      expect(ws.send).toHaveBeenCalledWith(
+        expect.stringContaining('"type":"error"')
+      );
+      expect(ws.send).toHaveBeenCalledWith(
+        expect.stringContaining('Session not found')
+      );
+      expect(ws.close).toHaveBeenCalledWith(4004, 'Session not found');
+    });
+
+    it('有效 session 注册连接并发送 connected 消息', () => {
+      const ws = createMockWSContext();
+      mockExists.mockReturnValue(true);
+
+      handleWebSocket(ws as any, 'session-1');
+
+      expect(ws.close).not.toHaveBeenCalled();
+      expect(ws.send).toHaveBeenCalledWith(
+        expect.stringContaining('"type":"connected"')
+      );
+      expect(mockUpdateStatus).toHaveBeenCalledWith('session-1', 'active');
+    });
+
+    it('连接被添加到 connections map', () => {
+      const ws = createMockWSContext();
+      mockExists.mockReturnValue(true);
+
+      handleWebSocket(ws as any, 'session-1');
+
+      const connections = getSessionConnections('session-1');
+      expect(connections.has(ws as any)).toBe(true);
+    });
+  });
+
+  describe('handleWebSocketMessage - 消息处理', () => {
+    beforeEach(() => {
+      mockExists.mockReturnValue(true);
+    });
+
+    it('处理 message 类型消息', async () => {
+      const ws = createMockWSContext();
+      const message = JSON.stringify({
+        type: 'message',
+        payload: { content: 'Hello AI' },
+      });
+
+      await handleWebSocketMessage(ws as any, 'session-1', message);
+
+      expect(mockAddMessage).toHaveBeenCalledWith('session-1', {
+        role: 'user',
+        content: 'Hello AI',
+      });
+      expect(processMessage).toHaveBeenCalledWith('session-1', 'Hello AI');
+    });
+
+    it('处理 cancel 类型消息', async () => {
+      const ws = createMockWSContext();
+      handleWebSocket(ws as any, 'session-1');
+
+      const message = JSON.stringify({ type: 'cancel' });
+      await handleWebSocketMessage(ws as any, 'session-1', message);
+
+      expect(cancelProcessing).toHaveBeenCalledWith('session-1');
+      // 应该广播 cancelled 消息
+      expect(ws.send).toHaveBeenCalledWith(
+        expect.stringContaining('"type":"cancelled"')
+      );
+    });
+
+    it('处理 permission_response 类型消息', async () => {
+      const ws = createMockWSContext();
+      const message = JSON.stringify({
+        type: 'permission_response',
+        payload: { requestId: 'req-123', allow: true, remember: false },
+      });
+
+      await handleWebSocketMessage(ws as any, 'session-1', message);
+
+      expect(handlePermissionResponse).toHaveBeenCalledWith('req-123', true, false);
+    });
+
+    it('未知消息类型返回错误', async () => {
+      const ws = createMockWSContext();
+      const message = JSON.stringify({ type: 'unknown_type' });
+
+      await handleWebSocketMessage(ws as any, 'session-1', message);
+
+      expect(ws.send).toHaveBeenCalledWith(
+        expect.stringContaining('"type":"error"')
+      );
+      expect(ws.send).toHaveBeenCalledWith(
+        expect.stringContaining('Unknown message type')
+      );
+    });
+
+    it('无效 JSON 返回错误', async () => {
+      const ws = createMockWSContext();
+      await handleWebSocketMessage(ws as any, 'session-1', 'invalid json');
+
+      expect(ws.send).toHaveBeenCalledWith(
+        expect.stringContaining('"type":"error"')
+      );
+    });
+
+    it('处理 ArrayBuffer 数据', async () => {
+      const ws = createMockWSContext();
+
+      const message = { type: 'message', payload: { content: 'ArrayBuffer test' } };
+      const encoder = new TextEncoder();
+      const buffer = encoder.encode(JSON.stringify(message)).buffer;
+
+      await handleWebSocketMessage(ws as any, 'session-1', buffer);
+
+      expect(mockAddMessage).toHaveBeenCalledWith('session-1', {
+        role: 'user',
+        content: 'ArrayBuffer test',
+      });
+    });
+
+    it('空 content 处理正确', async () => {
+      const ws = createMockWSContext();
+      const message = JSON.stringify({
+        type: 'message',
+        payload: {},
+      });
+
+      await handleWebSocketMessage(ws as any, 'session-1', message);
+
+      expect(mockAddMessage).toHaveBeenCalledWith('session-1', {
+        role: 'user',
+        content: '',
+      });
+    });
+  });
+
+  describe('handleWebSocketClose - 关闭处理', () => {
+    it('从 connections 中移除连接', () => {
+      const sessionId = getUniqueSessionId('close');
+      const ws = createMockWSContext();
+      mockExists.mockReturnValue(true);
+
+      handleWebSocket(ws as any, sessionId);
+      expect(getSessionConnections(sessionId).size).toBe(1);
+
+      handleWebSocketClose(ws as any, sessionId);
+      expect(getSessionConnections(sessionId).size).toBe(0);
+    });
+
+    it('最后一个连接关闭时更新 session 状态为 idle', () => {
+      const sessionId = getUniqueSessionId('idle');
+      const ws = createMockWSContext();
+      mockExists.mockReturnValue(true);
+
+      handleWebSocket(ws as any, sessionId);
+      handleWebSocketClose(ws as any, sessionId);
+
+      expect(mockUpdateStatus).toHaveBeenLastCalledWith(sessionId, 'idle');
+    });
+
+    it('多个连接时关闭一个不影响其他', () => {
+      const sessionId = getUniqueSessionId('multi');
+      const ws1 = createMockWSContext();
+      const ws2 = createMockWSContext();
+      mockExists.mockReturnValue(true);
+
+      handleWebSocket(ws1 as any, sessionId);
+      handleWebSocket(ws2 as any, sessionId);
+      expect(getSessionConnections(sessionId).size).toBe(2);
+
+      handleWebSocketClose(ws1 as any, sessionId);
+      expect(getSessionConnections(sessionId).size).toBe(1);
+      expect(getSessionConnections(sessionId).has(ws2 as any)).toBe(true);
+    });
+  });
+
+  describe('broadcastToSession - 会话广播', () => {
+    it('向所有连接发送消息', () => {
+      const ws1 = createMockWSContext();
+      const ws2 = createMockWSContext();
+      mockExists.mockReturnValue(true);
+
+      handleWebSocket(ws1 as any, 'session-1');
+      handleWebSocket(ws2 as any, 'session-1');
+
+      vi.clearAllMocks();
+
+      broadcastToSession('session-1', {
+        type: 'chunk',
+        sessionId: 'session-1',
+        payload: { content: 'test' },
+      });
+
+      expect(ws1.send).toHaveBeenCalledTimes(1);
+      expect(ws2.send).toHaveBeenCalledTimes(1);
+      expect(ws1.send).toHaveBeenCalledWith(expect.stringContaining('"type":"chunk"'));
+    });
+
+    it('无连接时不抛出错误', () => {
+      expect(() => {
+        broadcastToSession('non-existent', {
+          type: 'chunk',
+          sessionId: 'non-existent',
+          payload: {},
+        });
+      }).not.toThrow();
+    });
+
+    it('发送失败时继续发送给其他连接', () => {
+      mockExists.mockReturnValue(true);
+
+      const ws1 = createMockWSContext();
+      const ws2 = createMockWSContext();
+
+      handleWebSocket(ws1 as any, 'session-1');
+      handleWebSocket(ws2 as any, 'session-1');
+
+      // Now set up ws1.send to throw
+      ws1.send.mockImplementation(() => {
+        throw new Error('Connection closed');
+      });
+
+      expect(() => {
+        broadcastToSession('session-1', {
+          type: 'chunk',
+          sessionId: 'session-1',
+          payload: {},
+        });
+      }).not.toThrow();
+
+      expect(ws2.send).toHaveBeenCalled();
+    });
+  });
+
+  describe('getSessionConnections - 获取会话连接', () => {
+    it('返回会话的所有连接', () => {
+      const sessionId = getUniqueSessionId('get-conns');
+      const ws = createMockWSContext();
+      mockExists.mockReturnValue(true);
+
+      handleWebSocket(ws as any, sessionId);
+
+      const connections = getSessionConnections(sessionId);
+      expect(connections.size).toBe(1);
+    });
+
+    it('不存在的会话返回空 Set', () => {
+      const connections = getSessionConnections('non-existent-unique-12345');
+      expect(connections.size).toBe(0);
+    });
+  });
+
+  describe('getConnectionStats - 连接统计', () => {
+    it('返回正确的统计信息', () => {
+      // Get initial stats to account for connections from previous tests
+      const initialStats = getConnectionStats();
+
+      const sessionId1 = getUniqueSessionId('stats-1');
+      const sessionId2 = getUniqueSessionId('stats-2');
+      const ws1 = createMockWSContext();
+      const ws2 = createMockWSContext();
+      const ws3 = createMockWSContext();
+
+      mockExists.mockReturnValue(true);
+
+      handleWebSocket(ws1 as any, sessionId1);
+      handleWebSocket(ws2 as any, sessionId1);
+      handleWebSocket(ws3 as any, sessionId2);
+
+      const stats = getConnectionStats();
+      // We added 2 sessions and 3 connections
+      expect(stats.sessions).toBe(initialStats.sessions + 2);
+      expect(stats.connections).toBe(initialStats.connections + 3);
+    });
+  });
+});