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', () => {
  it('should provide WebSocket context to children', () => {
    const mockUseWebSocket = vi.spyOn(useWebSocketModule, 'useWebSocket');
    mockUseWebSocket.mockReturnValue({
      isConnected: true,
      socket: 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', () => {
    const mockUseWebSocket = vi.spyOn(useWebSocketModule, 'useWebSocket');
    mockUseWebSocket.mockReturnValue({
      isConnected: false,
      socket: 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', () => {
    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', () => {
    const mockUseWebSocket = vi.spyOn(useWebSocketModule, 'useWebSocket');
    
    // Initially disconnected
    mockUseWebSocket.mockReturnValue({
      isConnected: false,
      socket: 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,
    });

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

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