lvmh/web/src/ws.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { createWsManager } from "./ws";
import { FakeWebSocket } from "./test/setup";

const URL = "ws://srv/ws?token=t";

describe("createWsManager", () => {
	beforeEach(() => {
		vi.useFakeTimers();
	});

	afterEach(() => {
		vi.useRealTimers();
	});

	it("subscribe/unsubscribe/resubscribeAll are silent while the socket is not open", () => {
		const m = createWsManager(URL);
		const sock = FakeWebSocket.last(); // still CONNECTING
		m.subscribe("s1");
		m.unsubscribe("s1");
		m.resubscribeAll();
		expect(sock.sent).toEqual([]);
		sock.serverClose();
		m.resubscribeAll();
		expect(sock.sent).toEqual([]);
		m.close();
	});

	it("frames received after close() are ignored (handlers nulled)", () => {
		const m = createWsManager(URL);
		const sock = FakeWebSocket.last();
		sock.serverOpen();
		m.close();
		expect(() =>
			sock.serverMessage({ type: "session_list", sessions: [] }),
		).not.toThrow();
		expect(FakeWebSocket.instances.length).toBe(1);
	});

	it("onState receives transitions", () => {
		const states: string[] = [];
		const m = createWsManager(URL);
		const off = m.onState((s) => states.push(s));
		expect(states).toEqual(["connecting"]);
		expect(FakeWebSocket.last().url).toBe(URL);

		FakeWebSocket.last().serverOpen();
		expect(states).toEqual(["connecting", "open"]);

		off();
		FakeWebSocket.last().serverClose();
		expect(states).toEqual(["connecting", "open"]);
		m.close();
	});

	it("accepts url as function", () => {
		const m = createWsManager(() => URL);
		expect(FakeWebSocket.last().url).toBe(URL);
		m.close();
	});

	it("onFrame receives parsed frames; malformed frames dropped", () => {
		const m = createWsManager(URL);
		const frames: unknown[] = [];
		m.onFrame((f) => frames.push(f));
		const sock = FakeWebSocket.last();
		sock.serverOpen();

		sock.serverMessage({ type: "session_list", sessions: [] });
		expect(frames).toEqual([{ type: "session_list", sessions: [] }]);

		sock.onmessage?.({ data: "{bad json" } as MessageEvent);
		sock.onmessage?.({ data: JSON.stringify(null) } as MessageEvent);
		sock.onmessage?.({
			data: JSON.stringify({ noType: true }),
		} as MessageEvent);
		sock.onmessage?.({ data: JSON.stringify(42) } as MessageEvent);
		expect(frames).toHaveLength(1);

		m.close();
	});

	it("subscribe/unsubscribe send frames while open, resubscribe on reopen", () => {
		const m = createWsManager(URL);
		const sock = FakeWebSocket.last();
		sock.serverOpen();

		m.subscribe("s1");
		expect(sock.sent).toEqual([
			JSON.stringify({ type: "subscribe", sessionId: "s1" }),
		]);
		m.subscribe("s2");
		m.unsubscribe("s2");
		expect(sock.sent).toHaveLength(3);

		sock.serverClose();
		// silent drop: reconnect timer scheduled with jitter
		vi.spyOn(Math, "random").mockReturnValue(1);
		vi.advanceTimersByTime(500);
		const sock2 = FakeWebSocket.last();
		sock2.serverOpen();
		expect(sock2.sent).toEqual([
			JSON.stringify({ type: "subscribe", sessionId: "s1" }),
		]);
		expect(m.subscriptions()).toEqual(new Set(["s1"]));
		m.close();
	});

	it("resubscribeAll replays live subscriptions over an open socket", () => {
		const m = createWsManager(URL);
		const sock = FakeWebSocket.last();
		sock.serverOpen();
		m.subscribe("a");
		m.subscribe("b");
		sock.sent.length = 0;
		m.resubscribeAll();
		expect(sock.sent).toEqual([
			JSON.stringify({ type: "subscribe", sessionId: "a" }),
			JSON.stringify({ type: "subscribe", sessionId: "b" }),
		]);
		m.close();
	});

	it("full backoff ladder reaches the 30s cap", () => {
		vi.spyOn(Math, "random").mockReturnValue(1);
		const m = createWsManager(URL);
		const ladder = [500, 1000, 2000, 4000, 8000, 16000, 30000, 30000];
		let elapsed = 0;
		for (const delay of ladder) {
			FakeWebSocket.last().serverClose();
			elapsed += delay;
			vi.advanceTimersByTime(delay);
			FakeWebSocket.last().serverOpen();
		}
		expect(FakeWebSocket.instances.length).toBe(ladder.length + 1);
		void elapsed;
		m.close();
	});

	it("successful open resets the backoff attempt counter", () => {
		vi.spyOn(Math, "random").mockReturnValue(1);
		const m = createWsManager(URL);
		FakeWebSocket.last().serverClose();
		vi.advanceTimersByTime(500);
		FakeWebSocket.last().serverOpen();
		FakeWebSocket.last().serverClose();
		vi.advanceTimersByTime(500); // attempt reset to 0 by open
		expect(FakeWebSocket.instances.length).toBe(3);
		m.close();
	});

	it("auth failure (close 1008) fires onAuthError and never reconnects", () => {
		const m = createWsManager(URL);
		let authErrors = 0;
		m.onAuthError(() => {
			authErrors += 1;
		});
		const sock = FakeWebSocket.last();
		sock.serverOpen();
		sock.serverClose(1008);
		expect(authErrors).toBe(1);
		vi.advanceTimersByTime(120_000);
		expect(FakeWebSocket.instances.length).toBe(1);
		expect(m.subscriptions() instanceof Set).toBe(true);
		m.close();
	});

	it("close() tears the socket down and stops reconnection", () => {
		vi.spyOn(Math, "random").mockReturnValue(1);
		const m = createWsManager(URL);
		const sock = FakeWebSocket.last();
		sock.serverOpen();
		m.subscribe("s1");
		m.close();
		expect(sock.closeCode).toBe(4900);
		expect(sock.onclose).toBeNull();
		expect(sock.onmessage).toBeNull();
		vi.advanceTimersByTime(120_000);
		expect(FakeWebSocket.instances.length).toBe(1);
	});

	it("callbacks unsubscribe via returned off functions", () => {
		const m = createWsManager(URL);
		const offFrame = m.onFrame(() => undefined);
		const offAuth = m.onAuthError(() => undefined);
		const offState = m.onState(() => undefined);
		offFrame();
		offAuth();
		offState();
		const sock = FakeWebSocket.last();
		sock.serverOpen();
		sock.serverMessage({ type: "spawn_status", jobs: [] });
		sock.serverClose(1008);
		m.close();
		expect(true).toBe(true);
	});

	it("non-string ws data does not throw", () => {
		const m = createWsManager(URL);
		const sock = FakeWebSocket.last();
		sock.serverOpen();
		expect(() =>
			sock.onmessage?.({ data: {} } as unknown as MessageEvent),
		).not.toThrow();
		m.close();
	});
});