headquarter/apps/web/src/hooks/use-terminal-connection.test.ts

import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { renderHook, act } from "@testing-library/react";
import { useTerminalConnection } from "./use-terminal-connection";

class MockWebSocket {
	static instances: MockWebSocket[] = [];
	readyState: number = WebSocket.CONNECTING;
	onopen: ((ev: Event) => void) | null = null;
	onclose: ((ev: CloseEvent) => void) | null = null;
	onmessage: ((ev: MessageEvent) => void) | null = null;
	onerror: ((ev: Event) => void) | null = null;
	sent: (string | ArrayBuffer | Blob)[] = [];
	url = "";

	constructor(url: string) {
		this.url = url;
		MockWebSocket.instances.push(this);
	}

	send(data: string | ArrayBuffer | Blob) {
		this.sent.push(data);
	}

	close(code?: number, reason?: string) {
		this.readyState = WebSocket.CLOSED;
		if (this.onclose) {
			this.onclose(new CloseEvent("close", { code: code ?? 1000, reason }));
		}
	}

	simulateOpen() {
		this.readyState = WebSocket.OPEN;
		if (this.onopen) this.onopen(new Event("open"));
	}

	simulateMessage(data: string | ArrayBuffer | Blob) {
		if (this.onmessage) {
			this.onmessage(new MessageEvent("message", { data }));
		}
	}

	simulateError() {
		if (this.onerror) this.onerror(new Event("error"));
	}
}

describe("useTerminalConnection", () => {
	let originalWebSocket: typeof WebSocket;

	beforeEach(() => {
		originalWebSocket = globalThis.WebSocket;
		globalThis.WebSocket = MockWebSocket as unknown as typeof WebSocket;
		MockWebSocket.instances = [];
		vi.useFakeTimers();
		vi.stubGlobal("import", { meta: { env: { VITE_API_BASE_URL: "" } } });
	});

	afterEach(() => {
		globalThis.WebSocket = originalWebSocket;
		MockWebSocket.instances = [];
		vi.useRealTimers();
		vi.unstubAllGlobals();
	});

	it("starts in connecting state", () => {
		const { result } = renderHook(() =>
			useTerminalConnection({ instanceId: "inst-1" }),
		);
		expect(result.current.state.status).toBe("connecting");
		expect(MockWebSocket.instances).toHaveLength(1);
	});

	it("transitions to connected on websocket open", async () => {
		const { result } = renderHook(() =>
			useTerminalConnection({ instanceId: "inst-1" }),
		);

		act(() => {
			MockWebSocket.instances[0].simulateOpen();
		});

		expect(result.current.state.status).toBe("connected");
	});

	it("sends ping after interval", async () => {
		renderHook(() => useTerminalConnection({ instanceId: "inst-1" }));

		act(() => {
			MockWebSocket.instances[0].simulateOpen();
		});

		act(() => {
			vi.advanceTimersByTime(15000);
		});

		const pings = MockWebSocket.instances[0].sent.filter((s) =>
			typeof s === "string" ? s.includes("ping") : false,
		);
		expect(pings.length).toBeGreaterThanOrEqual(1);
	});

	it("handles pong and updates latency", async () => {
		const { result } = renderHook(() =>
			useTerminalConnection({ instanceId: "inst-1" }),
		);

		act(() => {
			MockWebSocket.instances[0].simulateOpen();
		});

		act(() => {
			vi.advanceTimersByTime(15000);
		});

		act(() => {
			MockWebSocket.instances[0].simulateMessage(
				JSON.stringify({ type: "pong", id: 1 }),
			);
		});

		expect(result.current.state.latency).not.toBeNull();
		expect(result.current.state.latency).toBeGreaterThanOrEqual(0);
	});

	it("reconnects with exponential backoff on close", async () => {
		const { result } = renderHook(() =>
			useTerminalConnection({ instanceId: "inst-1" }),
		);

		act(() => {
			MockWebSocket.instances[0].simulateOpen();
		});

		act(() => {
			MockWebSocket.instances[0].close(1006, "Abnormal closure");
		});

		expect(result.current.state.status).toBe("reconnecting");
		expect(result.current.state.attempt).toBe(1);

		act(() => {
			vi.advanceTimersByTime(1000);
		});

		expect(MockWebSocket.instances).toHaveLength(2);
	});

	it("max reconnect attempts leads to disconnected", async () => {
		const { result } = renderHook(() =>
			useTerminalConnection({ instanceId: "inst-1" }),
		);

		act(() => {
			MockWebSocket.instances[0].simulateOpen();
		});

		for (let i = 0; i < 11; i++) {
			const ws = MockWebSocket.instances[MockWebSocket.instances.length - 1];
			act(() => {
				ws.close(1006, "Abnormal closure");
			});
			const delay = Math.min(1000 * 2 ** i, 30000);
			act(() => {
				vi.advanceTimersByTime(delay);
			});
		}

		expect(result.current.state.status).toBe("disconnected");
		expect(result.current.state.error).toContain("Max reconnection");
	}, 30000);

	it("sends resize message with debounce", async () => {
		const { result } = renderHook(() =>
			useTerminalConnection({ instanceId: "inst-1" }),
		);

		act(() => {
			MockWebSocket.instances[0].simulateOpen();
		});

		act(() => {
			result.current.sendResize(120, 40);
		});

		// Before debounce
		expect(
			MockWebSocket.instances[0].sent.filter((s) =>
				typeof s === "string" ? s.includes("resize") : false,
			),
		).toHaveLength(0);

		act(() => {
			vi.advanceTimersByTime(250);
		});

		const resizes = MockWebSocket.instances[0].sent.filter((s) =>
			typeof s === "string" ? s.includes("resize") : false,
		);
		expect(resizes.length).toBeGreaterThanOrEqual(1);
	});

	it("throttles resize messages", async () => {
		const { result } = renderHook(() =>
			useTerminalConnection({ instanceId: "inst-1" }),
		);

		act(() => {
			MockWebSocket.instances[0].simulateOpen();
		});

		act(() => {
			result.current.sendResize(100, 30);
		});
		act(() => {
			vi.advanceTimersByTime(250);
		});

		act(() => {
			result.current.sendResize(101, 31);
		});
		act(() => {
			vi.advanceTimersByTime(250);
		});

		const resizes = MockWebSocket.instances[0].sent.filter((s) =>
			typeof s === "string" ? s.includes("resize") : false,
		);
		// Second resize throttled (within 500ms)
		expect(resizes.length).toBe(1);
	});

	it("sendInput sends data over websocket", async () => {
		const { result } = renderHook(() =>
			useTerminalConnection({ instanceId: "inst-1" }),
		);

		act(() => {
			MockWebSocket.instances[0].simulateOpen();
		});

		act(() => {
			result.current.sendInput("a");
		});

		expect(MockWebSocket.instances[0].sent).toContain("a");
	});

	it("triggers manual reconnect on reconnect()", async () => {
		const { result } = renderHook(() =>
			useTerminalConnection({ instanceId: "inst-1" }),
		);

		act(() => {
			MockWebSocket.instances[0].simulateOpen();
		});

		act(() => {
			result.current.reconnect();
		});

		expect(MockWebSocket.instances).toHaveLength(2);
	});

	it("calls onData callback with binary data", async () => {
		const onData = vi.fn();
		renderHook(() => useTerminalConnection({ instanceId: "inst-1", onData }));

		act(() => {
			MockWebSocket.instances[0].simulateOpen();
		});

		const buffer = new ArrayBuffer(3);
		act(() => {
			MockWebSocket.instances[0].simulateMessage(buffer);
		});

		expect(onData).toHaveBeenCalledWith(expect.any(Uint8Array));
	});

	it("calls onControl callback with control messages", async () => {
		const onControl = vi.fn();
		renderHook(() =>
			useTerminalConnection({ instanceId: "inst-1", onControl }),
		);

		act(() => {
			MockWebSocket.instances[0].simulateOpen();
		});

		act(() => {
			MockWebSocket.instances[0].simulateMessage(
				JSON.stringify({ type: "set_echo_state", enabled: false }),
			);
		});

		expect(onControl).toHaveBeenCalledWith(
			expect.objectContaining({ type: "set_echo_state", enabled: false }),
		);
	});

	it("serializes and restores scrollback", async () => {
		const serializeFn = vi.fn(() => "scrollback-content");
		const onRestoreScrollback = vi.fn();

		renderHook(
			() =>
				useTerminalConnection({
					instanceId: "inst-1",
					serializeFn,
					onRestoreScrollback,
				}),
			{ initialProps: {} },
		);

		act(() => {
			MockWebSocket.instances[0].simulateOpen();
		});

		// Disconnect
		act(() => {
			MockWebSocket.instances[0].close(1006, "gone");
		});

		expect(serializeFn).toHaveBeenCalled();

		act(() => {
			vi.advanceTimersByTime(1000);
		});

		// New connection opens
		act(() => {
			MockWebSocket.instances[
				MockWebSocket.instances.length - 1
			].simulateOpen();
		});

		expect(onRestoreScrollback).toHaveBeenCalledWith("scrollback-content");
	});
});