Playwright Socket.io WebSocket Mock

playwright-socketio-mock.ts

/**
 * Playwright Socket.io WebSocket Mock
 * 
 * A reusable utility for mocking Socket.io WebSocket connections in Playwright E2E tests.
 * Fully mocks the connection without hitting a real server - handles Engine.IO/Socket.io
 * protocol handshakes so the browser client thinks it's connected.
 * 
 * Usage:
 *   const mock = createWebSocketMock();
 *   await mock.setup(page, /your-ws-pattern/);
 *   mock.emit('yourEvent', { any: 'data' });
 *   mock.cleanup();
 * 
 * @see https://playwright.dev/docs/api/class-websocketroute
 */

import { Page, WebSocketRoute } from '@playwright/test';

class WebSocketMock {
  private wsRoute: WebSocketRoute | null = null;
  private pingInterval: ReturnType<typeof setInterval> | null = null;

  /**
   * Sets up a mocked WebSocket that handles Socket.io/Engine.IO protocol.
   * Call this BEFORE navigating to the page.
   *
   * @param page - Playwright page
   * @param urlPattern - WebSocket URL pattern to intercept
   */
  async setup(page: Page, urlPattern: RegExp) {
    await page.routeWebSocket(urlPattern, (ws) => {
      this.wsRoute = ws;

      // Engine.IO OPEN packet - tells client the connection is established
      ws.send(
        `0${JSON.stringify({
          sid: `mock-${Date.now()}`,
          upgrades: [],
          pingInterval: 25000,
          pingTimeout: 20000,
        })}`
      );

      ws.onMessage((message) => {
        if (typeof message === 'string') {
          // Engine.IO PING (2) -> respond with PONG (3)
          if (message === '2') {
            ws.send('3');
            return;
          }
          // Socket.io CONNECT (40) -> respond with CONNECT ACK
          if (message === '40') {
            ws.send('40{"sid":"mock-socket"}');
            return;
          }
        }
      });

      // Keep connection alive with periodic pings
      this.pingInterval = setInterval(() => {
        try {
          ws.send('2');
        } catch {
          // Connection closed
        }
      }, 20000);

      ws.onClose(() => {
        this.cleanup();
      });
    });
  }

  /**
   * Emits a Socket.io event to the browser.
   *
   * @param event - Event name
   * @param data - Event payload (will be JSON stringified)
   * @param stringifyData - If true (default), data is double-stringified for apps that JSON.parse the payload
   */
  emit(event: string, data: unknown, stringifyData = true) {
    if (!this.wsRoute) {
      throw new Error('WebSocket not set up. Call setup() first.');
    }
    // Socket.io protocol: 4 = MESSAGE, 2 = EVENT → "42"
    const payload = stringifyData ? JSON.stringify(JSON.stringify(data)) : JSON.stringify(data);
    const message = `42["${event}",${payload}]`;
    this.wsRoute.send(message);
  }

  /**
   * Cleans up the mock. Call this in afterEach or finally blocks.
   */
  cleanup() {
    if (this.pingInterval) {
      clearInterval(this.pingInterval);
      this.pingInterval = null;
    }
    this.wsRoute = null;
  }

  /**
   * Check if the WebSocket is set up and ready.
   */
  get isReady(): boolean {
    return this.wsRoute !== null;
  }
}

// Singleton for simple usage
let defaultMock: WebSocketMock | null = null;

export const getWebSocketMock = (): WebSocketMock => {
  if (!defaultMock) {
    defaultMock = new WebSocketMock();
  }
  return defaultMock;
};

export const createWebSocketMock = (): WebSocketMock => {
  return new WebSocketMock();
};

export const setupWebSocketMock = async (page: Page, urlPattern: RegExp): Promise<WebSocketMock> => {
  const mock = getWebSocketMock();
  await mock.setup(page, urlPattern);
  return mock;
};

export const emitWebSocketEvent = (event: string, data: unknown) => {
  getWebSocketMock().emit(event, data);
};

export const cleanupWebSocketMock = () => {
  if (defaultMock) {
    defaultMock.cleanup();
    defaultMock = null;
  }
};