stack/apps/web/src/providers/WebSocketProvider.test.tsx

import { describe, it, expect, vi } from "vitest";
import { render, screen } from "@testing-library/react";
import { WebSocketProvider, useWebSocketContext } from "./WebSocketProvider";
import * as useWebSocketModule from "../hooks/useWebSocket";

// Mock the useWebSocket hook
vi.mock("../hooks/useWebSocket");

describe("WebSocketProvider", (): void => {
  it("should provide WebSocket context to children", (): void => {
    const mockUseWebSocket = vi.spyOn(useWebSocketModule, "useWebSocket");
    mockUseWebSocket.mockReturnValue({
      isConnected: true,
      socket: null,
      connectionError: null,
    });

    function TestComponent(): React.JSX.Element {
      const { isConnected } = useWebSocketContext();
      return <div>{isConnected ? "Connected" : "Disconnected"}</div>;
    }

    render(
      <WebSocketProvider workspaceId="workspace-123" token="auth-token">
        <TestComponent />
      </WebSocketProvider>
    );

    expect(screen.getByText("Connected")).toBeInTheDocument();
  });

  it("should pass callbacks to useWebSocket hook", (): void => {
    const mockUseWebSocket = vi.spyOn(useWebSocketModule, "useWebSocket");
    mockUseWebSocket.mockReturnValue({
      isConnected: false,
      socket: null,
      connectionError: null,
    });

    const onTaskCreated = vi.fn();
    const onTaskUpdated = vi.fn();
    const onTaskDeleted = vi.fn();

    render(
      <WebSocketProvider
        workspaceId="workspace-123"
        token="auth-token"
        onTaskCreated={onTaskCreated}
        onTaskUpdated={onTaskUpdated}
        onTaskDeleted={onTaskDeleted}
      >
        <div>Test</div>
      </WebSocketProvider>
    );

    expect(mockUseWebSocket).toHaveBeenCalledWith("workspace-123", "auth-token", {
      onTaskCreated,
      onTaskUpdated,
      onTaskDeleted,
      onEventCreated: undefined,
      onEventUpdated: undefined,
      onEventDeleted: undefined,
      onProjectUpdated: undefined,
    });
  });

  it("should throw error when useWebSocketContext is used outside provider", (): void => {
    function TestComponent(): React.JSX.Element {
      useWebSocketContext();
      return <div>Test</div>;
    }

    // Suppress console.error for this test
    const originalError = console.error;
    console.error = vi.fn();

    expect(() => {
      render(<TestComponent />);
    }).toThrow("useWebSocketContext must be used within WebSocketProvider");

    console.error = originalError;
  });

  it("should update context when connection status changes", (): void => {
    const mockUseWebSocket = vi.spyOn(useWebSocketModule, "useWebSocket");

    // Initially disconnected
    mockUseWebSocket.mockReturnValue({
      isConnected: false,
      socket: null,
      connectionError: null,
    });

    function TestComponent(): React.JSX.Element {
      const { isConnected } = useWebSocketContext();
      return <div data-testid="status">{isConnected ? "Connected" : "Disconnected"}</div>;
    }

    const { rerender } = render(
      <WebSocketProvider workspaceId="workspace-123" token="auth-token">
        <TestComponent />
      </WebSocketProvider>
    );

    expect(screen.getByTestId("status")).toHaveTextContent("Disconnected");

    // Update to connected
    mockUseWebSocket.mockReturnValue({
      isConnected: true,
      socket: null,
      connectionError: null,
    });

    rerender(
      <WebSocketProvider workspaceId="workspace-123" token="auth-token">
        <TestComponent />
      </WebSocketProvider>
    );

    expect(screen.getByTestId("status")).toHaveTextContent("Connected");
  });
});