whale-town-end/test/utils/websocket-client.ts

/**
 * 原生 WebSocket 客户端测试工具
 * 
 * 用于替代 Socket.IO 客户端进行测试
 */

import WebSocket from 'ws';

export interface WebSocketTestClient {
  connect(): Promise<void>;
  disconnect(): void;
  send(event: string, data: any): void;
  on(event: string, callback: (data: any) => void): void;
  off(event: string, callback?: (data: any) => void): void;
  waitForEvent(event: string, timeout?: number): Promise<any>;
  isConnected(): boolean;
}

export class WebSocketTestClientImpl implements WebSocketTestClient {
  private ws: WebSocket | null = null;
  private eventHandlers = new Map<string, Set<(data: any) => void>>();
  private connected = false;

  constructor(private url: string) {}

  async connect(): Promise<void> {
    return new Promise((resolve, reject) => {
      this.ws = new WebSocket(this.url);

      this.ws.on('open', () => {
        this.connected = true;
        resolve();
      });

      this.ws.on('error', (error) => {
        reject(error);
      });

      this.ws.on('message', (data) => {
        try {
          const message = JSON.parse(data.toString());
          const { event, data: eventData } = message;
          
          const handlers = this.eventHandlers.get(event);
          if (handlers) {
            handlers.forEach(handler => handler(eventData));
          }
        } catch (error) {
          console.error('Failed to parse WebSocket message:', error);
        }
      });

      this.ws.on('close', () => {
        this.connected = false;
      });
    });
  }

  disconnect(): void {
    if (this.ws) {
      this.ws.close();
      this.ws = null;
      this.connected = false;
    }
  }

  send(event: string, data: any): void {
    if (this.ws && this.connected) {
      const message = JSON.stringify({ event, data });
      this.ws.send(message);
    } else {
      throw new Error('WebSocket is not connected');
    }
  }

  on(event: string, callback: (data: any) => void): void {
    if (!this.eventHandlers.has(event)) {
      this.eventHandlers.set(event, new Set());
    }
    this.eventHandlers.get(event)!.add(callback);
  }

  off(event: string, callback?: (data: any) => void): void {
    const handlers = this.eventHandlers.get(event);
    if (handlers) {
      if (callback) {
        handlers.delete(callback);
      } else {
        handlers.clear();
      }
    }
  }

  async waitForEvent(event: string, timeout: number = 5000): Promise<any> {
    return new Promise((resolve, reject) => {
      const timer = setTimeout(() => {
        this.off(event, handler);
        reject(new Error(`Timeout waiting for event: ${event}`));
      }, timeout);

      const handler = (data: any) => {
        clearTimeout(timer);
        this.off(event, handler);
        resolve(data);
      };

      this.on(event, handler);
    });
  }

  isConnected(): boolean {
    return this.connected;
  }
}

export function createWebSocketTestClient(url: string): WebSocketTestClient {
  return new WebSocketTestClientImpl(url);
}