stack/apps/web/src/hooks/useWebSocket.test.tsx

import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { renderHook, act, waitFor } from "@testing-library/react";
import { useWebSocket } from "./useWebSocket";
import type { Socket } from "socket.io-client";
import { io } from "socket.io-client";

// Mock socket.io-client
vi.mock("socket.io-client");

describe("useWebSocket", (): void => {
  let mockSocket: Partial<Socket>;
  let eventHandlers: Record<string, (data: unknown) => void>;

  beforeEach((): void => {
    eventHandlers = {};

    mockSocket = {
      on: vi.fn((event: string, handler: (...args: unknown[]) => void) => {
        eventHandlers[event] = handler;
        return mockSocket;
      }) as unknown as Socket["on"],
      off: vi.fn((event?: string) => {
        if (event && Object.hasOwn(eventHandlers, event)) {
          // eslint-disable-next-line @typescript-eslint/no-dynamic-delete
          delete eventHandlers[event];
        }
        return mockSocket;
      }) as unknown as Socket["off"],
      connect: vi.fn(),
      disconnect: vi.fn(),
      connected: false,
    };

    (io as unknown as ReturnType<typeof vi.fn>).mockReturnValue(mockSocket);
  });

  afterEach((): void => {
    vi.clearAllMocks();
  });

  it("should connect to WebSocket server on mount", (): void => {
    const workspaceId = "workspace-123";
    const token = "auth-token";

    renderHook(() => useWebSocket(workspaceId, token));

    expect(io).toHaveBeenCalledWith(expect.any(String), {
      auth: { token },
      query: { workspaceId },
      withCredentials: true,
    });
  });

  it("should disconnect on unmount", (): void => {
    const { unmount } = renderHook(() => useWebSocket("workspace-123", "token"));

    unmount();

    expect(mockSocket.disconnect).toHaveBeenCalled();
  });

  it("should update connection status on connect event", async (): Promise<void> => {
    mockSocket.connected = false;
    const { result } = renderHook(() => useWebSocket("workspace-123", "token"));

    expect(result.current.isConnected).toBe(false);

    act(() => {
      mockSocket.connected = true;
      eventHandlers.connect?.(undefined);
    });

    await waitFor(() => {
      expect(result.current.isConnected).toBe(true);
    });
  });

  it("should update connection status on disconnect event", async (): Promise<void> => {
    mockSocket.connected = true;
    const { result } = renderHook(() => useWebSocket("workspace-123", "token"));

    act(() => {
      eventHandlers.connect?.(undefined);
    });

    await waitFor(() => {
      expect(result.current.isConnected).toBe(true);
    });

    act(() => {
      mockSocket.connected = false;
      eventHandlers.disconnect?.(undefined);
    });

    await waitFor(() => {
      expect(result.current.isConnected).toBe(false);
    });
  });

  it("should handle task:created events", async (): Promise<void> => {
    const onTaskCreated = vi.fn();
    renderHook(() => useWebSocket("workspace-123", "token", { onTaskCreated }));

    const task = { id: "task-1", title: "New Task" };

    act(() => {
      eventHandlers["task:created"]?.(task);
    });

    await waitFor(() => {
      expect(onTaskCreated).toHaveBeenCalledWith(task);
    });
  });

  it("should handle task:updated events", async (): Promise<void> => {
    const onTaskUpdated = vi.fn();
    renderHook(() => useWebSocket("workspace-123", "token", { onTaskUpdated }));

    const task = { id: "task-1", title: "Updated Task" };

    act(() => {
      eventHandlers["task:updated"]?.(task);
    });

    await waitFor(() => {
      expect(onTaskUpdated).toHaveBeenCalledWith(task);
    });
  });

  it("should handle task:deleted events", async (): Promise<void> => {
    const onTaskDeleted = vi.fn();
    renderHook(() => useWebSocket("workspace-123", "token", { onTaskDeleted }));

    const payload = { id: "task-1" };

    act(() => {
      eventHandlers["task:deleted"]?.(payload);
    });

    await waitFor(() => {
      expect(onTaskDeleted).toHaveBeenCalledWith(payload);
    });
  });

  it("should handle event:created events", async (): Promise<void> => {
    const onEventCreated = vi.fn();
    renderHook(() => useWebSocket("workspace-123", "token", { onEventCreated }));

    const event = { id: "event-1", title: "New Event" };

    act(() => {
      eventHandlers["event:created"]?.(event);
    });

    await waitFor(() => {
      expect(onEventCreated).toHaveBeenCalledWith(event);
    });
  });

  it("should handle event:updated events", async (): Promise<void> => {
    const onEventUpdated = vi.fn();
    renderHook(() => useWebSocket("workspace-123", "token", { onEventUpdated }));

    const event = { id: "event-1", title: "Updated Event" };

    act(() => {
      eventHandlers["event:updated"]?.(event);
    });

    await waitFor(() => {
      expect(onEventUpdated).toHaveBeenCalledWith(event);
    });
  });

  it("should handle event:deleted events", async (): Promise<void> => {
    const onEventDeleted = vi.fn();
    renderHook(() => useWebSocket("workspace-123", "token", { onEventDeleted }));

    const payload = { id: "event-1" };

    act(() => {
      eventHandlers["event:deleted"]?.(payload);
    });

    await waitFor(() => {
      expect(onEventDeleted).toHaveBeenCalledWith(payload);
    });
  });

  it("should handle project:updated events", async (): Promise<void> => {
    const onProjectUpdated = vi.fn();
    renderHook(() => useWebSocket("workspace-123", "token", { onProjectUpdated }));

    const project = { id: "project-1", name: "Updated Project" };

    act(() => {
      eventHandlers["project:updated"]?.(project);
    });

    await waitFor(() => {
      expect(onProjectUpdated).toHaveBeenCalledWith(project);
    });
  });

  it("should reconnect with new workspace ID", (): void => {
    const { rerender } = renderHook(
      ({ workspaceId }: { workspaceId: string }) => useWebSocket(workspaceId, "token"),
      { initialProps: { workspaceId: "workspace-1" } }
    );

    expect(io).toHaveBeenCalledTimes(1);

    rerender({ workspaceId: "workspace-2" });

    expect(mockSocket.disconnect).toHaveBeenCalled();
    expect(io).toHaveBeenCalledTimes(2);
  });

  describe("stale closure prevention", (): void => {
    it("should NOT disconnect when callback functions change", (): void => {
      const onTaskCreated1 = vi.fn();
      const onTaskCreated2 = vi.fn();

      const { rerender } = renderHook(
        ({ onTaskCreated }: { onTaskCreated: (task: { id: string }) => void }) =>
          useWebSocket("workspace-123", "token", { onTaskCreated }),
        { initialProps: { onTaskCreated: onTaskCreated1 } }
      );

      expect(io).toHaveBeenCalledTimes(1);
      expect(mockSocket.disconnect).not.toHaveBeenCalled();

      // Change the callback - this should NOT cause a reconnect
      rerender({ onTaskCreated: onTaskCreated2 });

      // Socket should NOT have been disconnected or reconnected
      expect(mockSocket.disconnect).not.toHaveBeenCalled();
      expect(io).toHaveBeenCalledTimes(1);
    });

    it("should use the latest callback after callback change", async (): Promise<void> => {
      const onTaskCreated1 = vi.fn();
      const onTaskCreated2 = vi.fn();

      const { rerender } = renderHook(
        ({ onTaskCreated }: { onTaskCreated: (task: { id: string }) => void }) =>
          useWebSocket("workspace-123", "token", { onTaskCreated }),
        { initialProps: { onTaskCreated: onTaskCreated1 } }
      );

      // Emit event with first callback
      const task1 = { id: "task-1" };
      act(() => {
        eventHandlers["task:created"]?.(task1);
      });

      await waitFor(() => {
        expect(onTaskCreated1).toHaveBeenCalledWith(task1);
        expect(onTaskCreated2).not.toHaveBeenCalled();
      });

      // Update to new callback
      rerender({ onTaskCreated: onTaskCreated2 });

      // Emit another event - should use the new callback
      const task2 = { id: "task-2" };
      act(() => {
        eventHandlers["task:created"]?.(task2);
      });

      await waitFor(() => {
        expect(onTaskCreated2).toHaveBeenCalledWith(task2);
        // First callback should only have been called once (with task1)
        expect(onTaskCreated1).toHaveBeenCalledTimes(1);
      });
    });

    it("should NOT disconnect when multiple callbacks change simultaneously", (): void => {
      const callbacks1 = {
        onTaskCreated: vi.fn(),
        onTaskUpdated: vi.fn(),
        onEventCreated: vi.fn(),
      };
      const callbacks2 = {
        onTaskCreated: vi.fn(),
        onTaskUpdated: vi.fn(),
        onEventCreated: vi.fn(),
      };

      interface CallbackProps {
        onTaskCreated: (task: { id: string }) => void;
        onTaskUpdated: (task: { id: string }) => void;
        onEventCreated: (event: { id: string }) => void;
      }

      const { rerender } = renderHook(
        (props: CallbackProps) => useWebSocket("workspace-123", "token", props),
        { initialProps: callbacks1 }
      );

      expect(io).toHaveBeenCalledTimes(1);

      // Change all callbacks at once - should NOT cause reconnect
      rerender(callbacks2);

      expect(mockSocket.disconnect).not.toHaveBeenCalled();
      expect(io).toHaveBeenCalledTimes(1);
    });

    it("should handle callback being removed without reconnect", (): void => {
      const onTaskCreated = vi.fn();

      interface CallbackProps {
        onTaskCreated?: (task: { id: string }) => void;
      }

      const { rerender } = renderHook(
        (props: CallbackProps) => useWebSocket("workspace-123", "token", props),
        { initialProps: { onTaskCreated } as CallbackProps }
      );

      expect(io).toHaveBeenCalledTimes(1);

      // Remove callback - should NOT cause reconnect
      rerender({});

      expect(mockSocket.disconnect).not.toHaveBeenCalled();
      expect(io).toHaveBeenCalledTimes(1);
    });

    it("should handle callback being added without reconnect", (): void => {
      const onTaskCreated = vi.fn();

      interface CallbackProps {
        onTaskCreated?: (task: { id: string }) => void;
      }

      const { rerender } = renderHook(
        (props: CallbackProps) => useWebSocket("workspace-123", "token", props),
        { initialProps: {} as CallbackProps }
      );

      expect(io).toHaveBeenCalledTimes(1);

      // Add callback - should NOT cause reconnect
      rerender({ onTaskCreated });

      expect(mockSocket.disconnect).not.toHaveBeenCalled();
      expect(io).toHaveBeenCalledTimes(1);
    });
  });

  it("should clean up all event listeners on unmount", (): void => {
    const { unmount } = renderHook(() =>
      useWebSocket("workspace-123", "token", {
        onTaskCreated: vi.fn(),
        onTaskUpdated: vi.fn(),
        onTaskDeleted: vi.fn(),
      })
    );

    unmount();

    expect(mockSocket.off).toHaveBeenCalledWith("connect", expect.any(Function));
    expect(mockSocket.off).toHaveBeenCalledWith("disconnect", expect.any(Function));
    expect(mockSocket.off).toHaveBeenCalledWith("task:created", expect.any(Function));
    expect(mockSocket.off).toHaveBeenCalledWith("task:updated", expect.any(Function));
    expect(mockSocket.off).toHaveBeenCalledWith("task:deleted", expect.any(Function));
  });

  describe("connect_error handling", (): void => {
    it("should handle connect_error events and expose error state", async (): Promise<void> => {
      const { result } = renderHook(() => useWebSocket("workspace-123", "token"));

      expect(result.current.connectionError).toBeNull();

      const error = new Error("Connection refused");

      act(() => {
        eventHandlers.connect_error?.(error);
      });

      await waitFor(() => {
        expect(result.current.connectionError).toEqual({
          message: "Connection refused",
          type: "connect_error",
          description: "Failed to establish WebSocket connection",
        });
        expect(result.current.isConnected).toBe(false);
      });
    });

    it("should handle connect_error with missing message", async (): Promise<void> => {
      const { result } = renderHook(() => useWebSocket("workspace-123", "token"));

      const error = new Error();

      act(() => {
        eventHandlers.connect_error?.(error);
      });

      await waitFor(() => {
        expect(result.current.connectionError).toEqual({
          message: "Connection failed",
          type: "connect_error",
          description: "Failed to establish WebSocket connection",
        });
      });
    });

    it("should clear connection error on reconnect", async (): Promise<void> => {
      const { result, rerender } = renderHook(
        ({ workspaceId }: { workspaceId: string }) => useWebSocket(workspaceId, "token"),
        { initialProps: { workspaceId: "workspace-1" } }
      );

      // Simulate connect error
      act(() => {
        eventHandlers.connect_error?.(new Error("Connection failed"));
      });

      await waitFor(() => {
        expect(result.current.connectionError).not.toBeNull();
      });

      // Rerender with new workspace to trigger reconnect
      rerender({ workspaceId: "workspace-2" });

      // Connection error should be cleared when attempting new connection
      await waitFor(() => {
        expect(result.current.connectionError).toBeNull();
      });
    });

    it("should register connect_error handler on socket", (): void => {
      renderHook(() => useWebSocket("workspace-123", "token"));

      expect(mockSocket.on).toHaveBeenCalledWith("connect_error", expect.any(Function));
    });

    it("should clean up connect_error handler on unmount", (): void => {
      const { unmount } = renderHook(() => useWebSocket("workspace-123", "token"));

      unmount();

      expect(mockSocket.off).toHaveBeenCalledWith("connect_error", expect.any(Function));
    });
  });

  describe("WSS enforcement", (): void => {
    afterEach((): void => {
      vi.unstubAllEnvs();
      vi.resetModules();
    });

    it("should warn when using ws:// in production", async (): Promise<void> => {
      const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined);

      vi.stubEnv("NODE_ENV", "production");

      // Mock the config module to return insecure URL
      vi.doMock("@/lib/config", () => ({
        API_BASE_URL: "http://insecure-server.com",
      }));

      // Re-import to get fresh module with mocked config
      const { useWebSocket: useWebSocketMocked } = await import("./useWebSocket");

      renderHook(() => useWebSocketMocked("workspace-123", "token"));

      expect(consoleWarnSpy).toHaveBeenCalledWith(expect.stringContaining("[Security Warning]"));
      expect(consoleWarnSpy).toHaveBeenCalledWith(expect.stringContaining("insecure protocol"));

      consoleWarnSpy.mockRestore();
    });

    it("should not warn when using https:// in production", async (): Promise<void> => {
      const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined);

      vi.stubEnv("NODE_ENV", "production");

      // Mock the config module to return secure URL
      vi.doMock("@/lib/config", () => ({
        API_BASE_URL: "https://secure-server.com",
      }));

      // Re-import to get fresh module with mocked config
      const { useWebSocket: useWebSocketMocked } = await import("./useWebSocket");

      renderHook(() => useWebSocketMocked("workspace-123", "token"));

      expect(consoleWarnSpy).not.toHaveBeenCalled();

      consoleWarnSpy.mockRestore();
    });

    it("should not warn when using wss:// in production", async (): Promise<void> => {
      const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined);

      vi.stubEnv("NODE_ENV", "production");

      // Mock the config module to return secure WSS URL
      vi.doMock("@/lib/config", () => ({
        API_BASE_URL: "wss://secure-server.com",
      }));

      // Re-import to get fresh module with mocked config
      const { useWebSocket: useWebSocketMocked } = await import("./useWebSocket");

      renderHook(() => useWebSocketMocked("workspace-123", "token"));

      expect(consoleWarnSpy).not.toHaveBeenCalled();

      consoleWarnSpy.mockRestore();
    });

    it("should not warn in development mode even with http://", async (): Promise<void> => {
      const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined);

      vi.stubEnv("NODE_ENV", "development");

      // Mock the config module to return insecure URL (but we're in dev mode)
      vi.doMock("@/lib/config", () => ({
        API_BASE_URL: "http://localhost:3001",
      }));

      // Re-import to get fresh module with mocked config
      const { useWebSocket: useWebSocketMocked } = await import("./useWebSocket");

      renderHook(() => useWebSocketMocked("workspace-123", "token"));

      expect(consoleWarnSpy).not.toHaveBeenCalled();

      consoleWarnSpy.mockRestore();
    });
  });
});