stack/apps/web/src/spa/chat/test-support/fake-chat-socket.ts

import { vi } from 'vitest';
import type { ClientToServerEvents, ServerToClientEvents } from '@/lib/chat-contract';

type ServerEvent = keyof ServerToClientEvents;
type ClientEvent = keyof ClientToServerEvents;
type ServerHandler<K extends ServerEvent> = ServerToClientEvents[K];
type ClientPayload<K extends ClientEvent> = Parameters<ClientToServerEvents[K]>[0];

export interface EmittedEvent<K extends ClientEvent = ClientEvent> {
  event: K;
  payload: ClientPayload<K>;
}

/** The subset of a Socket.IO `ChatSocket` that `useChatConnection` drives. */
export interface FakeChatSocket {
  connected: boolean;
  connect(): FakeChatSocket;
  on<K extends ServerEvent>(event: K, handler: ServerHandler<K>): FakeChatSocket;
  off<K extends ServerEvent>(event: K, handler: ServerHandler<K>): FakeChatSocket;
  emit<K extends ClientEvent>(event: K, payload: ClientPayload<K>): FakeChatSocket;
}

/**
 * A typed in-memory stand-in for `getSocket()`. Unlike a bare
 * `(event: string, payload: unknown) => void` mock, every public method here is
 * checked against the real `/chat` contract — a typo'd event name or a payload
 * missing a required field fails to compile instead of silently no-op'ing at
 * runtime.
 */
/** Socket.IO's built-in connection-state events. Not part of the app-level
 * ServerToClientEvents contract, but real sockets always support them and
 * `useChatConnection` registers a `disconnect` handler on the real socket. */
type LifecycleEvent = 'connect' | 'disconnect';

export function createFakeChatSocket(): {
  socket: FakeChatSocket;
  listeners: Map<ServerEvent, Set<(payload: never) => void>>;
  emitted: EmittedEvent[];
  serverEmit<K extends ServerEvent>(
    event: K,
    payload: Parameters<ServerToClientEvents[K]>[0],
  ): void;
  /** Escape hatch for malformed-payload tests: bypasses the compile-time
   * payload contract to simulate a genuinely untrusted runtime value from the
   * server, e.g. a `session:info` with a non-array `availableThinkingLevels`. */
  serverEmitRaw(event: ServerEvent, payload: unknown): void;
  /** Simulates a transient Socket.IO `disconnect` — fires any handler(s)
   * registered via `socket.on('disconnect', ...)` without clearing any
   * listeners, mirroring how a real reconnecting socket behaves. */
  simulateDisconnect(): void;
  /** Simulates socket.io-client's automatic reconnect of the *same*
   * instance after a transient disconnect: marks the socket connected again
   * and fires any handler(s) registered via `socket.on('connect', ...)`,
   * without clearing or replacing any listeners. */
  simulateReconnect(): void;
} {
  const listeners = new Map<ServerEvent, Set<(payload: never) => void>>();
  const emitted: EmittedEvent[] = [];

  // Internal storage is intentionally keyed loosely (the per-event handler shape
  // varies by K, which a single Map can't express); the generic signatures on the
  // exported `socket`/`serverEmit` above and below are what keep test call sites
  // type-checked against ServerToClientEvents/ClientToServerEvents.
  const socket = {
    connected: false,
    connect: vi.fn(function connect(this: void) {
      socket.connected = true;
      return socket;
    }),
    on: vi.fn(function on(this: void, event: ServerEvent, handler: (payload: never) => void) {
      if (!listeners.has(event)) listeners.set(event, new Set());
      listeners.get(event)?.add(handler);
      return socket;
    }),
    off: vi.fn(function off(this: void, event: ServerEvent, handler: (payload: never) => void) {
      listeners.get(event)?.delete(handler);
      return socket;
    }),
    emit: vi.fn(function emit(this: void, event: ClientEvent, payload: unknown) {
      emitted.push({ event, payload } as EmittedEvent);
      return socket;
    }),
  } as unknown as FakeChatSocket;

  function serverEmit<K extends ServerEvent>(
    event: K,
    payload: Parameters<ServerToClientEvents[K]>[0],
  ): void {
    for (const handler of listeners.get(event) ?? []) {
      (handler as (payload: Parameters<ServerToClientEvents[K]>[0]) => void)(payload);
    }
  }

  function serverEmitRaw(event: ServerEvent, payload: unknown): void {
    for (const handler of listeners.get(event) ?? []) {
      (handler as (payload: unknown) => void)(payload);
    }
  }

  function simulateDisconnect(): void {
    socket.connected = false;
    const lifecycleKey = 'disconnect' satisfies LifecycleEvent as unknown as ServerEvent;
    for (const handler of listeners.get(lifecycleKey) ?? []) {
      (handler as () => void)();
    }
  }

  function simulateReconnect(): void {
    socket.connected = true;
    const lifecycleKey = 'connect' satisfies LifecycleEvent as unknown as ServerEvent;
    for (const handler of listeners.get(lifecycleKey) ?? []) {
      (handler as () => void)();
    }
  }

  return {
    socket,
    listeners,
    emitted,
    serverEmit,
    serverEmitRaw,
    simulateDisconnect,
    simulateReconnect,
  };
}