theo-agent-dashboard

ws.test.ts (8999B)

---

 import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
 import { WsClient } from "./ws";
 
 // Mock WebSocket
 class MockWebSocket {
   static instances: MockWebSocket[] = [];
   url: string;
   onopen: (() => void) | null = null;
   onclose: (() => void) | null = null;
   onmessage: ((event: { data: string }) => void) | null = null;
   onerror: (() => void) | null = null;
   sent: string[] = [];
   readyState = 1; // OPEN
 
   constructor(url: string) {
     this.url = url;
     MockWebSocket.instances.push(this);
   }
 
   send(data: string) {
     this.sent.push(data);
   }
 
   close() {
     this.readyState = 3;
   }
 
   simulateMessage(data: string) {
     this.onmessage?.({ data });
   }
 
   simulateOpen() {
     this.onopen?.();
   }
 
   simulateClose() {
     this.onclose?.();
   }
 
   simulateError() {
     this.onerror?.();
   }
 }
 
 beforeEach(() => {
   MockWebSocket.instances = [];
   vi.stubGlobal("WebSocket", MockWebSocket);
   vi.useFakeTimers();
 });
 
 afterEach(() => {
   vi.useRealTimers();
   vi.restoreAllMocks();
 });
 
 describe("WsClient", () => {
   it("connect creates a WebSocket", () => {
     const client = new WsClient("ws://localhost:8080/ws");
     client.connect();
 
     expect(MockWebSocket.instances).toHaveLength(1);
     expect(MockWebSocket.instances[0].url).toBe("ws://localhost:8080/ws");
   });
 
   it("send serializes JSON", () => {
     const client = new WsClient("ws://localhost:8080/ws");
     client.connect();
 
     const ws = MockWebSocket.instances[0];
     ws.simulateOpen();
 
     client.send("chat.message", { text: "hello" });
 
     expect(ws.sent).toHaveLength(1);
     expect(JSON.parse(ws.sent[0])).toEqual({
       type: "chat.message",
       data: { text: "hello" },
     });
   });
 
   it("send without data sends type only", () => {
     const client = new WsClient("ws://localhost:8080/ws");
     client.connect();
 
     const ws = MockWebSocket.instances[0];
     ws.simulateOpen();
 
     client.send("ping");
 
     expect(ws.sent).toHaveLength(1);
     expect(JSON.parse(ws.sent[0])).toEqual({ type: "ping", data: undefined });
   });
 
   it("send does nothing when not connected (readyState != 1)", () => {
     const client = new WsClient("ws://localhost:8080/ws");
     client.connect();
 
     const ws = MockWebSocket.instances[0];
     // readyState is 1 by default, but let's test when it's not open
     ws.readyState = 0; // CONNECTING
 
     client.send("chat.message", { text: "hello" });
     expect(ws.sent).toHaveLength(0);
   });
 
   it("receive dispatches to handler", () => {
     const handler = vi.fn();
     const client = new WsClient("ws://localhost:8080/ws");
     client.onMessage(handler);
     client.connect();
 
     const ws = MockWebSocket.instances[0];
     ws.simulateOpen();
     ws.simulateMessage(
       JSON.stringify({ type: "assistant.delta", data: { text: "hi" } }),
     );
 
     expect(handler).toHaveBeenCalledWith({
       type: "assistant.delta",
       data: { text: "hi" },
     });
   });
 
   it("receive dispatches to multiple handlers", () => {
     const handler1 = vi.fn();
     const handler2 = vi.fn();
     const client = new WsClient("ws://localhost:8080/ws");
     client.onMessage(handler1);
     client.onMessage(handler2);
     client.connect();
 
     const ws = MockWebSocket.instances[0];
     ws.simulateMessage(JSON.stringify({ type: "test" }));
 
     expect(handler1).toHaveBeenCalled();
     expect(handler2).toHaveBeenCalled();
   });
 
   it("ignores malformed JSON messages", () => {
     const handler = vi.fn();
     const client = new WsClient("ws://localhost:8080/ws");
     client.onMessage(handler);
     client.connect();
 
     const ws = MockWebSocket.instances[0];
     ws.simulateMessage("not json {{{");
 
     expect(handler).not.toHaveBeenCalled();
   });
 
   it("reconnects with exponential backoff on disconnect", () => {
     const client = new WsClient("ws://localhost:8080/ws", {
       maxReconnectAttempts: 5,
     });
     client.connect();
 
     const ws1 = MockWebSocket.instances[0];
     ws1.simulateOpen();
     ws1.simulateClose();
 
     // No immediate reconnect
     expect(MockWebSocket.instances).toHaveLength(1);
 
     // First reconnect after 1s
     vi.advanceTimersByTime(1000);
     expect(MockWebSocket.instances).toHaveLength(2);
 
     // Second reconnect after 2s
     const ws2 = MockWebSocket.instances[1];
     ws2.simulateOpen();
     ws2.simulateClose();
     vi.advanceTimersByTime(2000);
     expect(MockWebSocket.instances).toHaveLength(3);
   });
 
   it("stops reconnecting after max attempts", () => {
     const client = new WsClient("ws://localhost:8080/ws", {
       maxReconnectAttempts: 2,
     });
     client.connect();
 
     // Connect attempt 1: open then close
     MockWebSocket.instances[0].simulateOpen();
     MockWebSocket.instances[0].simulateClose();
     vi.advanceTimersByTime(1000);
 
     // 1st reconnect: instance 2 created
     expect(MockWebSocket.instances).toHaveLength(2);
     MockWebSocket.instances[1].simulateClose();
     vi.advanceTimersByTime(2000);
 
     // 2nd reconnect: instance 3 created
     expect(MockWebSocket.instances).toHaveLength(3);
 
     // No more reconnects after max (2) reached
     MockWebSocket.instances[2].simulateClose();
     vi.advanceTimersByTime(4000);
     expect(MockWebSocket.instances).toHaveLength(3);
   });
 
   it("caps reconnect delay at maxReconnectDelay", () => {
     const client = new WsClient("ws://localhost:8080/ws", {
       maxReconnectAttempts: 10,
       baseReconnectDelay: 1000,
       maxReconnectDelay: 5000,
     });
     client.connect();
 
     // Simulate many reconnects to reach the cap
     for (let i = 0; i < 4; i++) {
       MockWebSocket.instances[i].simulateClose();
       vi.advanceTimersByTime(5000); // enough time for any delay
     }
 
     // Should have created instances: 1 (initial) + 4 reconnects = 5
     expect(MockWebSocket.instances).toHaveLength(5);
   });
 
   it("destroy stops reconnection", () => {
     const client = new WsClient("ws://localhost:8080/ws");
     client.connect();
 
     MockWebSocket.instances[0].simulateOpen();
     client.destroy();
 
     MockWebSocket.instances[0].simulateClose();
     vi.advanceTimersByTime(5000);
 
     // No reconnect after destroy
     expect(MockWebSocket.instances).toHaveLength(1);
   });
 
   it("destroy clears reconnect timer", () => {
     const client = new WsClient("ws://localhost:8080/ws");
     client.connect();
 
     // Close to schedule reconnect
     MockWebSocket.instances[0].simulateOpen();
     MockWebSocket.instances[0].simulateClose();
 
     // Destroy before timer fires
     client.destroy();
 
     // Wait past the reconnect delay
     vi.advanceTimersByTime(5000);
 
     // No new instance created
     expect(MockWebSocket.instances).toHaveLength(1);
   });
 
   it("destroy closes websocket", () => {
     const client = new WsClient("ws://localhost:8080/ws");
     client.connect();
 
     const ws = MockWebSocket.instances[0];
     client.destroy();
 
     // After destroy, the ws should have been closed (readyState = 3)
     expect(ws.readyState).toBe(3);
   });
 
   it("onopen resets reconnectAttempts", () => {
     const client = new WsClient("ws://localhost:8080/ws", {
       maxReconnectAttempts: 5,
     });
     client.connect();
 
     // Simulate disconnect and reconnect
     MockWebSocket.instances[0].simulateOpen();
     MockWebSocket.instances[0].simulateClose();
     vi.advanceTimersByTime(1000);
 
     // Second connection opens
     MockWebSocket.instances[1].simulateOpen();
 
     // Disconnect again
     MockWebSocket.instances[1].simulateClose();
 
     // Should use 1s delay (reset) not 2s
     vi.advanceTimersByTime(1000);
     expect(MockWebSocket.instances).toHaveLength(3);
   });
 
   it("does not reconnect when destroyed flag is set in scheduleReconnect", () => {
     const client = new WsClient("ws://localhost:8080/ws", {
       maxReconnectAttempts: 5,
     });
     client.connect();
 
     MockWebSocket.instances[0].simulateOpen();
     MockWebSocket.instances[0].simulateClose();
 
     // Schedule is pending
     client.destroy();
 
     vi.advanceTimersByTime(1000);
     // Timer fires but destroyed is true, so no new connect
     expect(MockWebSocket.instances).toHaveLength(1);
   });
 
   it("onerror handler exists but does nothing", () => {
     const client = new WsClient("ws://localhost:8080/ws");
     client.connect();
 
     const ws = MockWebSocket.instances[0];
     // onerror should not throw
     expect(() => ws.simulateError()).not.toThrow();
   });
 
   it("defaults constructor options", () => {
     const client = new WsClient("ws://localhost:8080/ws");
     client.connect();
 
     // Should use default maxReconnectAttempts of 10
     // Simulate 10 failed reconnects
     for (let i = 0; i < 10; i++) {
       MockWebSocket.instances[i].simulateClose();
       vi.advanceTimersByTime(30000); // maxReconnectDelay
     }
 
     // 11th attempt should not happen
     const countBefore = MockWebSocket.instances.length;
     MockWebSocket.instances[countBefore - 1].simulateClose();
     vi.advanceTimersByTime(30000);
     expect(MockWebSocket.instances).toHaveLength(countBefore);
   });
 });