languard-servers-manager/frontend/src/__tests__/useWebSocket.test.tsx

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<typeof import("@tanstack/react-query")>();
  return {
    ...actual,
    useQueryClient: () => ({
      invalidateQueries: mockInvalidateQueries,
    }),
  };
});

function createWrapper() {
  const queryClient = new QueryClient();
  return function Wrapper({ children }: { children: ReactNode }) {
    return (
      <QueryClientProvider client={queryClient}>
        {children}
      </QueryClientProvider>
    );
  };
}

describe("useWebSocket", () => {
  let mockWsInstance: {
    close: ReturnType<typeof vi.fn>;
    onopen: ((ev: Event) => void) | null;
    onmessage: ((ev: MessageEvent) => void) | null;
    onclose: ((ev: CloseEvent) => void) | null;
    onerror: ((ev: Event) => void) | null;
  };

  let MockWebSocket: ReturnType<typeof vi.fn>;

  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<typeof useAuthStore>);
  });

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

  it("should not connect when no token exists", () => {
    vi.mocked(useAuthStore).mockReturnValue({
      token: null,
    } as unknown as ReturnType<typeof useAuthStore>);

    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();
  });
});