Refactor relay transport socket tests (#882)

packages/server/src/server/relay-transport.test.ts

@@ -1,124 +1,115 @@
 import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
-
-const wsMock = vi.hoisted(() => {
-  class MockWebSocket {
-    static readonly CONNECTING = 0;
-    static readonly OPEN = 1;
-    static readonly CLOSING = 2;
-    static readonly CLOSED = 3;
-    static instances: MockWebSocket[] = [];
-
-    readonly url: string;
-    readonly options: unknown;
-    readyState = MockWebSocket.CONNECTING;
-    sent: string[] = [];
-    terminateCalls = 0;
-    private listeners = new Map<string, Array<(...args: unknown[]) => void>>();
-
-    constructor(url: string, options?: unknown) {
-      this.url = url;
-      this.options = options;
-      MockWebSocket.instances.push(this);
-    }
-
-    static reset() {
-      MockWebSocket.instances = [];
-    }
-
-    on(event: string, listener: (...args: unknown[]) => void) {
-      const handlers = this.listeners.get(event) ?? [];
-      handlers.push(listener);
-      this.listeners.set(event, handlers);
-      return this;
-    }
-
-    once(event: string, listener: (...args: unknown[]) => void) {
-      const wrapped = (...args: unknown[]) => {
-        this.off(event, wrapped);
-        listener(...args);
-      };
-      return this.on(event, wrapped);
-    }
-
-    close(code?: number, reason?: string) {
-      this.readyState = MockWebSocket.CLOSED;
-      this.emit("close", code ?? 1000, reason ?? "");
-    }
-
-    terminate() {
-      this.terminateCalls += 1;
-      this.readyState = MockWebSocket.CLOSED;
-      this.emit("close", 1006, "");
-    }
-
-    send(data: string) {
-      if (this.readyState !== MockWebSocket.OPEN) {
-        throw new Error(`WebSocket not open (readyState=${this.readyState})`);
-      }
-      this.sent.push(data);
-    }
-
-    open() {
-      this.readyState = MockWebSocket.OPEN;
-      this.emit("open");
-    }
-
-    message(data: unknown) {
-      this.emit("message", data);
-    }
-
-    error(err: unknown) {
-      this.emit("error", err);
-    }
-
-    private off(event: string, listener: (...args: unknown[]) => void) {
-      const handlers = this.listeners.get(event) ?? [];
-      this.listeners.set(
-        event,
-        handlers.filter((handler) => handler !== listener),
-      );
-    }
-
-    private emit(event: string, ...args: unknown[]) {
-      const handlers = this.listeners.get(event) ?? [];
-      for (const handler of handlers.slice()) {
-        handler(...args);
-      }
-    }
-  }
-
-  return { MockWebSocket };
-});
-
-vi.mock("ws", () => ({
-  default: wsMock.MockWebSocket,
-  WebSocket: wsMock.MockWebSocket,
-}));
-
 import type pino from "pino";
 import { startRelayTransport } from "./relay-transport";
 
 function createMockLogger() {
+  const messages: { level: "debug" | "info" | "warn" | "error"; args: unknown[] }[] = [];
   const logger = {
-    child: vi.fn(() => logger),
+    messages,
-    debug: vi.fn(),
+    child: () => logger,
-    info: vi.fn(),
+    debug: (...args: unknown[]) => messages.push({ level: "debug", args }),
-    warn: vi.fn(),
+    info: (...args: unknown[]) => messages.push({ level: "info", args }),
-    error: vi.fn(),
+    warn: (...args: unknown[]) => messages.push({ level: "warn", args }),
+    error: (...args: unknown[]) => messages.push({ level: "error", args }),
   };
   return logger;
 }
 
-function hasLogMessage(mockFn: ReturnType<typeof vi.fn>, message: string): boolean {
+type TestLogger = ReturnType<typeof createMockLogger>;
-  return mockFn.mock.calls.some((call) => call.some((arg) => arg === message));
+
+function hasLogMessage(logger: TestLogger, level: "info" | "warn", message: string): boolean {
+  return logger.messages.some((entry) => {
+    return entry.level === level && entry.args.some((arg) => arg === message);
+  });
+}
+
+class FakeRelayWebSocket {
+  static readonly CONNECTING = 0;
+  static readonly OPEN = 1;
+  static readonly CLOSED = 3;
+
+  readyState = FakeRelayWebSocket.CONNECTING;
+  sent: Array<string | Uint8Array | ArrayBuffer> = [];
+  terminateCalls = 0;
+  private readonly listeners = new Map<string, Array<(...args: unknown[]) => void>>();
+
+  constructor(readonly url: string) {}
+
+  on(event: string, listener: (...args: unknown[]) => void) {
+    const handlers = this.listeners.get(event) ?? [];
+    handlers.push(listener);
+    this.listeners.set(event, handlers);
+  }
+
+  once(event: string, listener: (...args: unknown[]) => void) {
+    const wrapped = (...args: unknown[]) => {
+      this.off(event, wrapped);
+      listener(...args);
+    };
+    this.on(event, wrapped);
+  }
+
+  close(code?: number, reason?: string) {
+    this.readyState = FakeRelayWebSocket.CLOSED;
+    this.emit("close", code ?? 1000, reason ?? "");
+  }
+
+  terminate() {
+    this.terminateCalls += 1;
+    this.readyState = FakeRelayWebSocket.CLOSED;
+    this.emit("close", 1006, "");
+  }
+
+  send(data: string | Uint8Array | ArrayBuffer) {
+    if (this.readyState !== FakeRelayWebSocket.OPEN) {
+      throw new Error(`WebSocket not open (readyState=${this.readyState})`);
+    }
+    this.sent.push(data);
+  }
+
+  open() {
+    this.readyState = FakeRelayWebSocket.OPEN;
+    this.emit("open");
+  }
+
+  message(data: unknown) {
+    this.emit("message", data);
+  }
+
+  private off(event: string, listener: (...args: unknown[]) => void) {
+    const handlers = this.listeners.get(event) ?? [];
+    this.listeners.set(
+      event,
+      handlers.filter((handler) => handler !== listener),
+    );
+  }
+
+  private emit(event: string, ...args: unknown[]) {
+    const handlers = this.listeners.get(event) ?? [];
+    for (const handler of handlers.slice()) {
+      handler(...args);
+    }
+  }
+}
+
+function createFakeWebSockets() {
+  const sockets: FakeRelayWebSocket[] = [];
+  return {
+    sockets,
+    createWebSocket(url: string) {
+      const socket = new FakeRelayWebSocket(url);
+      sockets.push(socket);
+      return socket;
+    },
+  };
 }
 
 describe("relay-transport control lifecycle", () => {
   const controllers: Array<{ stop: () => Promise<void> }> = [];
-  const MockWebSocket = wsMock.MockWebSocket;
+  let relay: ReturnType<typeof createFakeWebSockets>;
 
   beforeEach(() => {
-    MockWebSocket.reset();
+    relay = createFakeWebSockets();
   });
 
   afterEach(async () => {
@@ -135,18 +126,19 @@ describe("relay-transport control lifecycle", () => {
       relayEndpoint: "relay.paseo.sh:443",
       relayUseTls: true,
       serverId: "srv_test",
+      createWebSocket: relay.createWebSocket,
     });
     controllers.push(controller);
 
-    const control = MockWebSocket.instances[0];
+    const control = relay.sockets[0];
     expect(control).toBeDefined();
 
     control.open();
-    expect(hasLogMessage(logger.info, "relay_control_connected")).toBe(false);
+    expect(hasLogMessage(logger, "info", "relay_control_connected")).toBe(false);
     expect(control.sent.length).toBeGreaterThan(0);
 
     control.message(JSON.stringify({ type: "pong", ts: Date.now() }));
-    expect(hasLogMessage(logger.info, "relay_control_connected")).toBe(true);
+    expect(hasLogMessage(logger, "info", "relay_control_connected")).toBe(true);
   });
 
   test("terminates and reconnects when control socket opens but never becomes ready", () => {
@@ -158,18 +150,19 @@ describe("relay-transport control lifecycle", () => {
       relayEndpoint: "relay.paseo.sh:443",
       relayUseTls: true,
       serverId: "srv_test",
+      createWebSocket: relay.createWebSocket,
     });
     controllers.push(controller);
 
-    const firstControl = MockWebSocket.instances[0];
+    const firstControl = relay.sockets[0];
     firstControl.open();
 
     vi.advanceTimersByTime(8_000);
-    expect(hasLogMessage(logger.warn, "relay_control_ready_timeout_terminating")).toBe(true);
+    expect(hasLogMessage(logger, "warn", "relay_control_ready_timeout_terminating")).toBe(true);
     expect(firstControl.terminateCalls).toBe(1);
 
     vi.advanceTimersByTime(1_000);
-    expect(MockWebSocket.instances.length).toBeGreaterThanOrEqual(2);
+    expect(relay.sockets.length).toBeGreaterThanOrEqual(2);
   });
 
   test("terminates stale control sockets in under one minute", () => {
@@ -181,47 +174,56 @@ describe("relay-transport control lifecycle", () => {
       relayEndpoint: "relay.paseo.sh:443",
       relayUseTls: true,
       serverId: "srv_test",
+      createWebSocket: relay.createWebSocket,
     });
     controllers.push(controller);
 
-    const control = MockWebSocket.instances[0];
+    const control = relay.sockets[0];
     control.open();
     control.message(JSON.stringify({ type: "pong", ts: Date.now() }));
-    logger.warn.mockClear();
+    logger.messages.length = 0;
 
     vi.advanceTimersByTime(40_000);
-    expect(hasLogMessage(logger.warn, "relay_control_stale_terminating")).toBe(true);
+    expect(hasLogMessage(logger, "warn", "relay_control_stale_terminating")).toBe(true);
     expect(control.terminateCalls).toBe(1);
   });
 
   test("passes stable relay external session metadata when attaching data socket", async () => {
     const logger = createMockLogger();
-    const attachSocket = vi.fn(async () => {});
+    const attachedSockets: unknown[] = [];
+    const attachedMetadata: unknown[] = [];
+    const attachSocket = async (socket: unknown, metadata: unknown) => {
+      attachedSockets.push(socket);
+      attachedMetadata.push(metadata);
+    };
     const controller = startRelayTransport({
       logger: logger as unknown as pino.Logger,
       attachSocket,
       relayEndpoint: "relay.paseo.sh:443",
       relayUseTls: true,
       serverId: "srv_test",
+      createWebSocket: relay.createWebSocket,
     });
     controllers.push(controller);
 
-    const control = MockWebSocket.instances[0];
+    const control = relay.sockets[0];
     control.open();
     control.message(JSON.stringify({ type: "pong", ts: Date.now() }));
     control.message(JSON.stringify({ type: "connected", connectionId: "clt_test" }));
 
-    const dataSocket = MockWebSocket.instances[1];
+    const dataSocket = relay.sockets[1];
     expect(dataSocket).toBeDefined();
     dataSocket.open();
 
     await Promise.resolve();
 
-    expect(attachSocket).toHaveBeenCalledTimes(1);
+    expect(attachedSockets).toEqual([dataSocket]);
-    expect(attachSocket).toHaveBeenCalledWith(dataSocket, {
+    expect(attachedMetadata).toEqual([
-      transport: "relay",
+      {
-      externalSessionKey: "session:clt_test",
+        transport: "relay",
-    });
+        externalSessionKey: "session:clt_test",
+      },
+    ]);
   });
 
   test("uses relayUseTls for control and data socket URLs", () => {
@@ -232,15 +234,16 @@ describe("relay-transport control lifecycle", () => {
       relayEndpoint: "[::1]:443",
       relayUseTls: true,
       serverId: "srv_test",
+      createWebSocket: relay.createWebSocket,
     });
     controllers.push(controller);
 
-    const control = MockWebSocket.instances[0];
+    const control = relay.sockets[0];
     control.open();
     control.message(JSON.stringify({ type: "pong", ts: Date.now() }));
     control.message(JSON.stringify({ type: "connected", connectionId: "clt_test" }));
 
-    expect(MockWebSocket.instances[0]?.url).toMatch(/^wss:\/\/\[::1\]\/ws\?/);
+    expect(relay.sockets[0]?.url).toMatch(/^wss:\/\/\[::1\]\/ws\?/);
-    expect(MockWebSocket.instances[1]?.url).toMatch(/^wss:\/\/\[::1\]\/ws\?/);
+    expect(relay.sockets[1]?.url).toMatch(/^wss:\/\/\[::1\]\/ws\?/);
   });
 });

packages/server/src/server/relay-transport.ts

@@ -18,6 +18,7 @@ interface RelayTransportOptions {
   relayUseTls: boolean;
   serverId: string;
   daemonKeyPair?: KeyPair;
+  createWebSocket?: RelayWebSocketFactory;
 }
 
 export interface RelayTransportController {
@@ -32,6 +33,16 @@ interface RelaySocketLike {
   once: (event: "close" | "error", listener: (...args: unknown[]) => void) => void;
 }
 
+interface RelayWebSocketLike extends RelaySocketLike {
+  terminate: () => void;
+  on: (
+    event: "open" | "message" | "close" | "error",
+    listener: (...args: unknown[]) => void,
+  ) => void;
+}
+
+type RelayWebSocketFactory = (url: string) => RelayWebSocketLike;
+
 type ControlMessage =
   | { type: "sync"; connectionIds: string[] }
   | { type: "connected"; connectionId: string }
@@ -42,6 +53,11 @@ type ControlMessage =
 const CONTROL_PING_INTERVAL_MS = 10_000;
 const CONTROL_STALE_TIMEOUT_MS = 30_000;
 const CONTROL_READY_TIMEOUT_MS = 8_000;
+const RELAY_WEBSOCKET_OPTIONS = { handshakeTimeout: 10_000, perMessageDeflate: false } as const;
+
+function createDefaultRelayWebSocket(url: string): RelayWebSocketLike {
+  return new WebSocket(url, RELAY_WEBSOCKET_OPTIONS);
+}
 
 function normalizeRelaySendPayload(data: string | Uint8Array | ArrayBuffer): string | ArrayBuffer {
   if (typeof data === "string") return data;
@@ -106,14 +122,15 @@ export function startRelayTransport({
   relayUseTls,
   serverId,
   daemonKeyPair,
+  createWebSocket = createDefaultRelayWebSocket,
 }: RelayTransportOptions): RelayTransportController {
   const relayLogger = logger.child({ module: "relay-transport" });
 
   let stopped = false;
-  let controlWs: WebSocket | null = null;
+  let controlWs: RelayWebSocketLike | null = null;
   let reconnectTimeout: ReturnType<typeof setTimeout> | null = null;
   let reconnectAttempt = 0;
-  const dataSockets = new Map<string, WebSocket>(); // connectionId -> ws
+  const dataSockets = new Map<string, RelayWebSocketLike>(); // connectionId -> ws
   let controlKeepaliveInterval: ReturnType<typeof setInterval> | null = null;
   let controlReadyTimeout: ReturnType<typeof setTimeout> | null = null;
   let controlLastSeenAt = 0;
@@ -161,7 +178,7 @@ export function startRelayTransport({
       serverId,
       role: "server",
     });
-    const socket = new WebSocket(url, { handshakeTimeout: 10_000, perMessageDeflate: false });
+    const socket = createWebSocket(url);
     controlWs = socket;
     let controlConnected = false;
 
@@ -338,7 +355,7 @@ export function startRelayTransport({
       role: "server",
       connectionId,
     });
-    const socket = new WebSocket(url, { handshakeTimeout: 10_000, perMessageDeflate: false });
+    const socket = createWebSocket(url);
     dataSockets.set(connectionId, socket);
 
     let attached = false;
@@ -397,7 +414,7 @@ export function startRelayTransport({
 }
 
 async function attachEncryptedSocket(
-  socket: WebSocket,
+  socket: RelayWebSocketLike,
   daemonKeyPair: KeyPair,
   logger: pino.Logger,
   attachSocket: (ws: RelaySocketLike, metadata?: ExternalSocketMetadata) => Promise<void>,
@@ -426,7 +443,10 @@ async function attachEncryptedSocket(
   }
 }
 
-function createRelayTransportAdapter(socket: WebSocket, logger: pino.Logger): RelayTransport {
+function createRelayTransportAdapter(
+  socket: RelayWebSocketLike,
+  logger: pino.Logger,
+): RelayTransport {
   const relayTransport: RelayTransport = {
     send: (data) => {
       try {
@@ -445,10 +465,11 @@ function createRelayTransportAdapter(socket: WebSocket, logger: pino.Logger): Re
   };
 
   socket.on("message", (data, isBinary) => {
-    relayTransport.onmessage?.(normalizeMessageData(data, isBinary));
+    relayTransport.onmessage?.(normalizeMessageData(data, isBinary === true));
   });
   socket.on("close", (code, reason) => {
-    relayTransport.onclose?.(code, reason.toString());
+    const closeCode = typeof code === "number" ? code : 1006;
+    relayTransport.onclose?.(closeCode, String(reason ?? ""));
   });
   socket.on("error", (err) => {
     relayTransport.onerror?.(err instanceof Error ? err : new Error(String(err)));