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 { io, Socket } from 'socket.io-client';

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

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

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

    mockSocket = {
      on: vi.fn((event: string, handler: (data: unknown) => void) => {
        eventHandlers[event] = handler;
        return mockSocket as Socket;
      }),
      off: vi.fn((event: string) => {
        delete eventHandlers[event];
        return mockSocket as Socket;
      }),
      connect: vi.fn(),
      disconnect: vi.fn(),
      connected: false,
    };

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

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

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

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

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

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

    unmount();

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

  it('should update connection status on connect event', async () => {
    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 () => {
    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 () => {
    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 () => {
    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 () => {
    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 () => {
    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 () => {
    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 () => {
    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 () => {
    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', () => {
    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);
  });

  it('should clean up all event listeners on unmount', () => {
    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));
  });
});