959 lines · 24.9 KB · TypeScript
RawDownload
import { WS_METHODS } from "@t3tools/contracts";import * as Duration from "effect/Duration";import * as Effect from "effect/Effect";import * as Stream from "effect/Stream";import { afterEach, beforeEach, describe, expect, it, vi } from "vite-plus/test"; import { WsTransport } from "./wsTransport.ts"; type WsEventType = "open" | "message" | "close" | "error";type WsEvent = { code?: number; data?: unknown; reason?: string; type?: string };type WsListener = (event?: WsEvent) => void; const sockets: MockWebSocket[] = []; class MockWebSocket {  static readonly CONNECTING = 0;  static readonly OPEN = 1;  static readonly CLOSING = 2;  static readonly CLOSED = 3;   readyState = MockWebSocket.CONNECTING;  readonly sent: string[] = [];  readonly url: string;  private readonly listeners = new Map<WsEventType, Set<WsListener>>();   constructor(url: string) {    this.url = url;    sockets.push(this);  }   addEventListener(type: WsEventType, listener: WsListener) {    const listeners = this.listeners.get(type) ?? new Set<WsListener>();    listeners.add(listener);    this.listeners.set(type, listeners);  }   removeEventListener(type: WsEventType, listener: WsListener) {    this.listeners.get(type)?.delete(listener);  }   send(data: string) {    this.sent.push(data);  }   close(code = 1000, reason = "") {    this.readyState = MockWebSocket.CLOSED;    this.emit("close", { code, reason, type: "close" });  }   open() {    this.readyState = MockWebSocket.OPEN;    this.emit("open", { type: "open" });  }   serverMessage(data: unknown) {    this.emit("message", { data, type: "message" });  }   error() {    this.emit("error", { type: "error" });  }   private emit(type: WsEventType, event?: WsEvent) {    const listeners = this.listeners.get(type);    if (!listeners) return;    for (const listener of listeners) {      listener(event);    }  }} const originalWebSocket = globalThis.WebSocket;const originalFetch = globalThis.fetch;const transports: WsTransport[] = []; function getSocket(): MockWebSocket {  const socket = sockets.at(-1);  if (!socket) {    throw new Error("Expected a websocket instance");  }  return socket;} async function waitFor(assertion: () => void, timeoutMs = 1_000): Promise<void> {  const startedAt = performance.now();  for (;;) {    try {      assertion();      return;    } catch (error) {      if (performance.now() - startedAt >= timeoutMs) {        throw error;      }      await Effect.runPromise(Effect.sleep(Duration.millis(10)));    }  }} function createTransport(...args: ConstructorParameters<typeof WsTransport>): WsTransport {  const transport = new WsTransport(...args);  transports.push(transport);  return transport;} beforeEach(() => {  vi.useRealTimers();  sockets.length = 0;  transports.length = 0;   Object.defineProperty(globalThis, "window", {    configurable: true,    value: {      location: {        origin: "http://localhost:3020",        hostname: "localhost",        port: "3020",        protocol: "http:",      },      desktopBridge: undefined,    },  });  Object.defineProperty(globalThis, "navigator", {    configurable: true,    value: { onLine: true },  });   globalThis.WebSocket = MockWebSocket as unknown as typeof WebSocket;}); afterEach(async () => {  await Promise.allSettled(transports.map((transport) => transport.dispose()));  transports.length = 0;  globalThis.WebSocket = originalWebSocket;  globalThis.fetch = originalFetch;  vi.restoreAllMocks();}); describe("WsTransport", () => {  it("normalizes root websocket urls to /ws and preserves query params", async () => {    const transport = createTransport("ws://localhost:3020/?token=secret-token");     await waitFor(() => {      expect(sockets).toHaveLength(1);    });     expect(getSocket().url).toBe("ws://localhost:3020/ws?token=secret-token");    await transport.dispose();  });   it("uses an explicit secure websocket base url", async () => {    const transport = createTransport("wss://app.example.com");     await waitFor(() => {      expect(sockets).toHaveLength(1);    });     expect(getSocket().url).toBe("wss://app.example.com/ws");    await transport.dispose();  });   it("uses an explicit insecure websocket base url for remote backends", async () => {    const transport = createTransport("ws://192.168.1.44:3773");     await waitFor(() => {      expect(sockets).toHaveLength(1);    });     expect(getSocket().url).toBe("ws://192.168.1.44:3773/ws");    await transport.dispose();  });   it("supports async websocket url providers", async () => {    const transport = createTransport(async () => "wss://remote.example.com/?wsTicket=dynamic");     await waitFor(() => {      expect(sockets).toHaveLength(1);    });     expect(getSocket().url).toBe("wss://remote.example.com/ws?wsTicket=dynamic");    await transport.dispose();  });   it("invokes optional lifecycle handlers when the socket opens and closes", async () => {    const onOpen = vi.fn();    const onClose = vi.fn();    const transport = createTransport("ws://localhost:3020", {      onOpen,      onClose,    });     await waitFor(() => {      expect(sockets).toHaveLength(1);    });     const socket = getSocket();    socket.open();     await waitFor(() => {      expect(onOpen).toHaveBeenCalledOnce();    });     socket.close(1012, "service restart");     await waitFor(() => {      expect(onClose).toHaveBeenCalledWith(        {          code: 1012,          reason: "service restart",        },        {          intentional: false,        },      );    });     await transport.dispose();  });   it("tracks heartbeat freshness from websocket pongs", async () => {    const nowSpy = vi.spyOn(performance, "now").mockReturnValue(1_000);    const onHeartbeatPong = vi.fn();    const transport = createTransport("ws://localhost:3020", { onHeartbeatPong });     await waitFor(() => {      expect(sockets).toHaveLength(1);    });     expect(transport.isHeartbeatFresh()).toBe(false);     const socket = getSocket();    socket.open();    socket.serverMessage(JSON.stringify({ _tag: "Pong" }));     await waitFor(() => {      expect(onHeartbeatPong).toHaveBeenCalledOnce();    });     expect(transport.isHeartbeatFresh()).toBe(true);    expect(transport.isHeartbeatFresh(500)).toBe(true);     nowSpy.mockReturnValue(1_501);    expect(transport.isHeartbeatFresh(500)).toBe(false);     await transport.dispose();  });   it("clears heartbeat freshness when reconnecting", async () => {    vi.spyOn(performance, "now").mockReturnValue(1_000);    const onHeartbeatPong = vi.fn();    const transport = createTransport("ws://localhost:3020", { onHeartbeatPong });     await waitFor(() => {      expect(sockets).toHaveLength(1);    });     const firstSocket = getSocket();    firstSocket.open();    firstSocket.serverMessage(JSON.stringify({ _tag: "Pong" }));     await waitFor(() => {      expect(onHeartbeatPong).toHaveBeenCalledOnce();    });    expect(transport.isHeartbeatFresh()).toBe(true);     await transport.reconnect();     expect(transport.isHeartbeatFresh()).toBe(false);     await transport.dispose();  });   it("does not report an intentional dispose as a close", async () => {    const onClose = vi.fn();    const transport = createTransport("ws://localhost:3020", { onClose });     await waitFor(() => {      expect(sockets).toHaveLength(1);    });     getSocket().open();    await transport.dispose();     expect(onClose).not.toHaveBeenCalled();  });   it("ignores stale socket lifecycle events after reconnect starts a new session", async () => {    const onClose = vi.fn();    const transport = createTransport("ws://localhost:3020", { onClose });     await waitFor(() => {      expect(sockets).toHaveLength(1);    });     const firstSocket = getSocket();    firstSocket.open();     await transport.reconnect();     await waitFor(() => {      expect(sockets).toHaveLength(2);    });     firstSocket.close(1006, "stale close");     expect(onClose).not.toHaveBeenCalled();     await transport.dispose();  });   it("reconnects the websocket session without disposing the transport", async () => {    const transport = createTransport("ws://localhost:3020");     await waitFor(() => {      expect(sockets).toHaveLength(1);    });     const firstSocket = getSocket();    firstSocket.open();     await transport.reconnect();     await waitFor(() => {      expect(sockets).toHaveLength(2);    });     const secondSocket = getSocket();    expect(secondSocket).not.toBe(firstSocket);    expect(firstSocket.readyState).toBe(MockWebSocket.CLOSED);     const requestPromise = transport.request((client) =>      client[WS_METHODS.serverUpsertKeybinding]({        command: "terminal.toggle",        key: "ctrl+k",      }),    );     secondSocket.open();     await waitFor(() => {      expect(secondSocket.sent).toHaveLength(1);    });     const requestMessage = JSON.parse(secondSocket.sent[0] ?? "{}") as { id: string };    secondSocket.serverMessage(      JSON.stringify({        _tag: "Exit",        requestId: requestMessage.id,        exit: {          _tag: "Success",          value: {            keybindings: [],            issues: [],          },        },      }),    );     await expect(requestPromise).resolves.toEqual({      keybindings: [],      issues: [],    });     await transport.dispose();  });   it("sends unary RPC requests and resolves successful exits", async () => {    const transport = createTransport("ws://localhost:3020");     const requestPromise = transport.request((client) =>      client[WS_METHODS.serverUpsertKeybinding]({        command: "terminal.toggle",        key: "ctrl+k",      }),    );     await waitFor(() => {      expect(sockets).toHaveLength(1);    });     const socket = getSocket();    socket.open();     await waitFor(() => {      expect(socket.sent).toHaveLength(1);    });     const requestMessage = JSON.parse(socket.sent[0] ?? "{}") as {      _tag: string;      id: string;      payload: unknown;      tag: string;    };    expect(requestMessage).toMatchObject({      _tag: "Request",      tag: WS_METHODS.serverUpsertKeybinding,      payload: {        command: "terminal.toggle",        key: "ctrl+k",      },    });     socket.serverMessage(      JSON.stringify({        _tag: "Exit",        requestId: requestMessage.id,        exit: {          _tag: "Success",          value: {            keybindings: [],            issues: [],          },        },      }),    );     await expect(requestPromise).resolves.toEqual({      keybindings: [],      issues: [],    });     await transport.dispose();  });   it("delivers stream chunks to subscribers", async () => {    const transport = createTransport("ws://localhost:3020");    const listener = vi.fn();     const unsubscribe = transport.subscribe(      (client) => client[WS_METHODS.subscribeServerLifecycle]({}),      listener,    );    await waitFor(() => {      expect(sockets).toHaveLength(1);    });     const socket = getSocket();    socket.open();     await waitFor(() => {      expect(socket.sent).toHaveLength(1);    });     const requestMessage = JSON.parse(socket.sent[0] ?? "{}") as { id: string; tag: string };    expect(requestMessage.tag).toBe(WS_METHODS.subscribeServerLifecycle);     const welcomeEvent = {      version: 1,      sequence: 1,      type: "welcome",      payload: {        environment: {          environmentId: "environment-local",          label: "Local environment",          platform: { os: "darwin", arch: "arm64" },          serverVersion: "0.0.0-test",          capabilities: { repositoryIdentity: true },        },        cwd: "/tmp/workspace",        projectName: "workspace",      },    };     socket.serverMessage(      JSON.stringify({        _tag: "Chunk",        requestId: requestMessage.id,        values: [welcomeEvent],      }),    );     await waitFor(() => {      expect(listener).toHaveBeenCalledWith(welcomeEvent);    });     unsubscribe();    await transport.dispose();  });   it("re-subscribes stream listeners after the stream exits", async () => {    const transport = createTransport("ws://localhost:3020");    const listener = vi.fn();    const onResubscribe = vi.fn();     const unsubscribe = transport.subscribe(      (client) => client[WS_METHODS.subscribeServerLifecycle]({}),      listener,      { onResubscribe },    );    await waitFor(() => {      expect(sockets).toHaveLength(1);    });     const socket = getSocket();    socket.open();     await waitFor(() => {      expect(socket.sent).toHaveLength(1);    });     const firstRequest = JSON.parse(socket.sent[0] ?? "{}") as { id: string };

Lines 1–500 of 959.