import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { renderHook, act } from "@testing-library/react"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import type { ReactNode } from "react"; import { useWebSocket } from "@/hooks/useWebSocket"; import { useAuthStore } from "@/store/auth.store"; vi.mock("@/store/auth.store", () => ({ useAuthStore: vi.fn(), })); const mockInvalidateQueries = vi.fn(); vi.mock("@tanstack/react-query", async (importOriginal) => { const actual = await importOriginal(); return { ...actual, useQueryClient: () => ({ invalidateQueries: mockInvalidateQueries, }), }; }); function createWrapper() { const queryClient = new QueryClient(); return function Wrapper({ children }: { children: ReactNode }) { return ( {children} ); }; } describe("useWebSocket", () => { let mockWsInstance: { close: ReturnType; onopen: ((ev: Event) => void) | null; onmessage: ((ev: MessageEvent) => void) | null; onclose: ((ev: CloseEvent) => void) | null; onerror: ((ev: Event) => void) | null; }; let MockWebSocket: ReturnType; beforeEach(() => { vi.useFakeTimers(); localStorage.clear(); mockInvalidateQueries.mockClear(); const handlers = { onopen: null as ((ev: Event) => void) | null, onmessage: null as ((ev: MessageEvent) => void) | null, onclose: null as ((ev: CloseEvent) => void) | null, onerror: null as ((ev: Event) => void) | null, }; mockWsInstance = { close: vi.fn(), get onopen() { return handlers.onopen; }, set onopen(v) { handlers.onopen = v; }, get onmessage() { return handlers.onmessage; }, set onmessage(v) { handlers.onmessage = v; }, get onclose() { return handlers.onclose; }, set onclose(v) { handlers.onclose = v; }, get onerror() { return handlers.onerror; }, set onerror(v) { handlers.onerror = v; }, }; MockWebSocket = vi.fn(function (this: typeof mockWsInstance, _url: string) { return mockWsInstance; }); vi.stubGlobal("WebSocket", MockWebSocket); vi.mocked(useAuthStore).mockReturnValue({ token: "test-token", } as unknown as ReturnType); }); afterEach(() => { vi.useRealTimers(); vi.restoreAllMocks(); }); it("should not connect when no token exists", () => { vi.mocked(useAuthStore).mockReturnValue({ token: null, } as unknown as ReturnType); renderHook(() => useWebSocket(), { wrapper: createWrapper() }); expect(MockWebSocket).not.toHaveBeenCalled(); }); it("should create WebSocket connection when token exists", () => { renderHook(() => useWebSocket(), { wrapper: createWrapper() }); expect(MockWebSocket).toHaveBeenCalled(); }); it("should include token in WebSocket URL", () => { renderHook(() => useWebSocket(), { wrapper: createWrapper() }); const calledUrl = MockWebSocket.mock.calls[0][0] as string; expect(calledUrl).toContain("token=test-token"); }); it("should invalidate queries on server_status event", () => { renderHook(() => useWebSocket(), { wrapper: createWrapper() }); act(() => { if (mockWsInstance.onmessage) { mockWsInstance.onmessage({ data: JSON.stringify({ type: "server_status", server_id: 1, data: { status: "running" }, }), } as unknown as MessageEvent); } }); expect(mockInvalidateQueries).toHaveBeenCalled(); }); it("should close WebSocket on unmount", () => { const { unmount } = renderHook(() => useWebSocket(), { wrapper: createWrapper(), }); unmount(); expect(mockWsInstance.close).toHaveBeenCalled(); }); });