flyer-crawler.projectium.com/src/hooks/useWebSocket.test.ts

// src/hooks/useWebSocket.test.ts
import { renderHook, act } from '@testing-library/react';
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { useWebSocket } from './useWebSocket';
import { eventBus } from '../services/eventBus';

// Mock eventBus
vi.mock('../services/eventBus', () => ({
  eventBus: {
    dispatch: vi.fn(),
  },
}));

// A mock WebSocket class for testing
class MockWebSocket {
  static instances: MockWebSocket[] = [];
  static CONNECTING = 0;
  static OPEN = 1;
  static CLOSING = 2;
  static CLOSED = 3;

  url: string;
  readyState: number;
  onopen: () => void = () => {};
  onclose: (event: { code: number; reason: string; wasClean: boolean }) => void = () => {};
  onmessage: (event: { data: string }) => void = () => {};
  onerror: (error: Event) => void = () => {};
  send = vi.fn();
  close = vi.fn((code = 1000, reason = 'Client disconnecting') => {
    if (this.readyState === MockWebSocket.CLOSED || this.readyState === MockWebSocket.CLOSING)
      return;
    this.readyState = MockWebSocket.CLOSING;
    setTimeout(() => {
      this.readyState = MockWebSocket.CLOSED;
      this.onclose({ code, reason, wasClean: code === 1000 });
    }, 0);
  });

  constructor(url: string) {
    this.url = url;
    this.readyState = MockWebSocket.CONNECTING;
    MockWebSocket.instances.push(this);
  }

  // --- Test Helper Methods ---
  _open() {
    this.readyState = MockWebSocket.OPEN;
    this.onopen();
  }

  _message(data: object | string) {
    if (typeof data === 'string') {
      this.onmessage({ data });
    } else {
      this.onmessage({ data: JSON.stringify(data) });
    }
  }

  _error(errorEvent = new Event('error')) {
    this.onerror(errorEvent);
  }

  _close(code: number, reason: string) {
    if (this.readyState === MockWebSocket.CLOSED) return;
    this.readyState = MockWebSocket.CLOSED;
    this.onclose({ code, reason, wasClean: code === 1000 });
  }

  static get lastInstance(): MockWebSocket | undefined {
    return this.instances[this.instances.length - 1];
  }

  static clearInstances() {
    this.instances = [];
  }
}

describe('useWebSocket Hook', () => {
  const mockToken = 'test-token';
  let consoleErrorSpy: ReturnType<typeof vi.spyOn>;

  beforeEach(() => {
    vi.useFakeTimers();
    global.WebSocket = MockWebSocket as any;
    MockWebSocket.clearInstances();

    Object.defineProperty(window, 'location', {
      value: { protocol: 'https:', host: 'testhost.com' },
      writable: true,
    });

    Object.defineProperty(document, 'cookie', {
      writable: true,
      value: `accessToken=${mockToken}`,
    });

    vi.clearAllMocks();
    consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
  });

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

  it('should not connect on mount if autoConnect is false', () => {
    renderHook(() => useWebSocket({ autoConnect: false }));
    expect(MockWebSocket.instances).toHaveLength(0);
  });

  it('should auto-connect on mount by default', () => {
    const { result } = renderHook(() => useWebSocket());
    expect(MockWebSocket.instances).toHaveLength(1);
    expect(MockWebSocket.lastInstance?.url).toBe(`wss://testhost.com/ws?token=${mockToken}`);
    expect(result.current.isConnecting).toBe(true);
  });

  it('should set an error state if no access token is found', () => {
    document.cookie = ''; // No token
    const { result } = renderHook(() => useWebSocket());

    expect(result.current.isConnected).toBe(false);
    expect(result.current.isConnecting).toBe(false);
    expect(result.current.error).toBe('No access token found. Please log in.');
    expect(MockWebSocket.instances).toHaveLength(0);
  });

  it('should transition to connected state on WebSocket open', () => {
    const onConnect = vi.fn();
    const { result } = renderHook(() => useWebSocket({ onConnect }));

    expect(result.current.isConnecting).toBe(true);
    act(() => MockWebSocket.lastInstance?._open());

    expect(result.current.isConnected).toBe(true);
    expect(result.current.isConnecting).toBe(false);
    expect(onConnect).toHaveBeenCalled();
  });

  it('should handle incoming messages and dispatch to eventBus', () => {
    renderHook(() => useWebSocket());
    act(() => MockWebSocket.lastInstance?._open());

    const dealData = { flyerId: 1 };
    act(() => MockWebSocket.lastInstance?._message({ type: 'deal-notification', data: dealData }));
    expect(eventBus.dispatch).toHaveBeenCalledWith('notification:deal', dealData);
  });

  it('should log an error for invalid JSON messages', () => {
    renderHook(() => useWebSocket());
    act(() => MockWebSocket.lastInstance?._open());

    const invalidJson = 'this is not json';
    act(() => MockWebSocket.lastInstance?._message(invalidJson));

    expect(consoleErrorSpy).toHaveBeenCalledWith(
      '[WebSocket] Failed to parse message:',
      expect.any(SyntaxError),
    );
    expect(eventBus.dispatch).not.toHaveBeenCalled();
  });

  it('should respond to ping with pong', () => {
    renderHook(() => useWebSocket());
    act(() => MockWebSocket.lastInstance?._open());

    act(() => MockWebSocket.lastInstance?._message({ type: 'ping', data: {} }));

    expect(MockWebSocket.lastInstance?.send).toHaveBeenCalledWith(
      expect.stringContaining('"type":"pong"'),
    );
  });

  it('should disconnect and clean up when disconnect is called', () => {
    const onDisconnect = vi.fn();
    const { result } = renderHook(() => useWebSocket({ onDisconnect }));
    act(() => MockWebSocket.lastInstance?._open());

    act(() => result.current.disconnect());

    expect(MockWebSocket.lastInstance?.close).toHaveBeenCalledWith(1000, 'Client disconnecting');
    expect(result.current.isConnected).toBe(false);

    act(() => vi.runAllTimers());
    expect(onDisconnect).toHaveBeenCalled();

    // Ensure no reconnection attempt is made
    expect(MockWebSocket.instances).toHaveLength(1);
  });

  it('should attempt to reconnect on unexpected close', () => {
    const { result } = renderHook(() => useWebSocket({ reconnectDelay: 1000 }));
    act(() => MockWebSocket.lastInstance?._open());

    act(() => MockWebSocket.lastInstance?._close(1006, 'Abnormal closure'));
    expect(result.current.isConnected).toBe(false);

    act(() => vi.advanceTimersByTime(1000));
    expect(MockWebSocket.instances).toHaveLength(2);
    expect(result.current.isConnecting).toBe(true);
  });

  it('should use exponential backoff for reconnection', () => {
    renderHook(() => useWebSocket({ reconnectDelay: 1000, maxReconnectAttempts: 3 }));
    act(() => MockWebSocket.lastInstance?._open());

    // 1st failure -> 1s delay
    act(() => MockWebSocket.lastInstance?._close(1006, 'Abnormal'));
    act(() => vi.advanceTimersByTime(1000));
    expect(MockWebSocket.instances).toHaveLength(2);

    // 2nd failure -> 2s delay
    act(() => MockWebSocket.lastInstance?._close(1006, 'Abnormal'));
    act(() => vi.advanceTimersByTime(2000));
    expect(MockWebSocket.instances).toHaveLength(3);

    // 3rd failure -> 4s delay
    act(() => MockWebSocket.lastInstance?._close(1006, 'Abnormal'));
    act(() => vi.advanceTimersByTime(4000));
    expect(MockWebSocket.instances).toHaveLength(4);
  });

  it('should stop reconnecting after maxReconnectAttempts', () => {
    const { result } = renderHook(() =>
      useWebSocket({ reconnectDelay: 100, maxReconnectAttempts: 1 }),
    );
    act(() => MockWebSocket.lastInstance?._open());

    // 1st failure
    act(() => MockWebSocket.lastInstance?._close(1006, 'Abnormal'));
    act(() => vi.advanceTimersByTime(100));
    expect(MockWebSocket.instances).toHaveLength(2);

    // 2nd failure (should be the last)
    act(() => MockWebSocket.lastInstance?._close(1006, 'Abnormal'));
    act(() => vi.advanceTimersByTime(5000));

    expect(MockWebSocket.instances).toHaveLength(2); // No new instance
    expect(result.current.error).toBe('Failed to reconnect after multiple attempts');
  });

  it('should reset reconnect attempts on a successful connection', () => {
    renderHook(() => useWebSocket({ reconnectDelay: 100, maxReconnectAttempts: 2 }));
    act(() => MockWebSocket.lastInstance?._open());

    act(() => MockWebSocket.lastInstance?._close(1006, 'Abnormal'));
    act(() => vi.advanceTimersByTime(100)); // 1st reconnect attempt
    expect(MockWebSocket.instances).toHaveLength(2);

    act(() => MockWebSocket.lastInstance?._open()); // Reconnect succeeds

    act(() => MockWebSocket.lastInstance?._close(1006, 'Abnormal'));
    act(() => vi.advanceTimersByTime(100)); // Delay should be reset to base
    expect(MockWebSocket.instances).toHaveLength(3);
  });

  it('should send a message when connected', () => {
    const { result } = renderHook(() => useWebSocket());
    act(() => MockWebSocket.lastInstance?._open());

    const message = { type: 'ping' as const, data: {}, timestamp: new Date().toISOString() };
    act(() => result.current.send(message));

    expect(MockWebSocket.lastInstance?.send).toHaveBeenCalledWith(JSON.stringify(message));
  });

  it('should warn when trying to send a message while not connected', () => {
    const consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
    const { result } = renderHook(() => useWebSocket());
    // Do not open connection

    const message = { type: 'ping' as const, data: {}, timestamp: new Date().toISOString() };
    act(() => result.current.send(message));

    expect(MockWebSocket.lastInstance?.send).not.toHaveBeenCalled();
    expect(consoleWarnSpy).toHaveBeenCalledWith('[WebSocket] Cannot send message: not connected');
    consoleWarnSpy.mockRestore();
  });

  it('should clean up on unmount', () => {
    const { unmount } = renderHook(() => useWebSocket());
    const instance = MockWebSocket.lastInstance;

    unmount();

    expect(instance?.close).toHaveBeenCalled();
    act(() => vi.advanceTimersByTime(5000));
    expect(MockWebSocket.instances).toHaveLength(1); // No new reconnect attempts
  });

  it('should maintain stable function references across rerenders', () => {
    const { result, rerender } = renderHook(() => useWebSocket());

    const initialConnect = result.current.connect;
    const initialDisconnect = result.current.disconnect;
    const initialSend = result.current.send;

    rerender();

    expect(result.current.connect).toBe(initialConnect);
    expect(result.current.disconnect).toBe(initialDisconnect);
    expect(result.current.send).toBe(initialSend);
  });
});