AetherForge/server/web/src/context/WebSocketProvider.test.tsx

/**
 * @vitest-environment happy-dom
 */
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { act, render, renderHook } from '@testing-library/react';
import { WebSocketProvider } from './WebSocketProvider';
import { useWebSocketContext } from './WebSocketContext';
import { useWebSocket } from '../hooks/useWebSocket';
import { setStoredAuth, clearStoredAuth } from '../api/auth';
import { mockAgent, mockShare } from '../test/fixtures';

type WSListener = ((event: { data: string }) => void) | null;

class MockWebSocket {
  static instances: MockWebSocket[] = [];
  static OPEN = 1;
  static CONNECTING = 0;
  static CLOSED = 3;

  url: string;
  readyState = MockWebSocket.CONNECTING;
  onopen: (() => void) | null = null;
  onclose: (() => void) | null = null;
  onerror: (() => void) | null = null;
  onmessage: WSListener = null;

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

  close() {
    this.readyState = MockWebSocket.CLOSED;
    this.onclose?.();
  }

  emitMessage(data: unknown) {
    this.onmessage?.({ data: JSON.stringify(data) });
  }

  emitOpen() {
    this.readyState = MockWebSocket.OPEN;
    this.onopen?.();
  }
}

describe('WebSocketProvider', () => {
  beforeEach(() => {
    sessionStorage.clear();
    MockWebSocket.instances = [];
    vi.stubGlobal('WebSocket', MockWebSocket as unknown as typeof WebSocket);
    Object.defineProperty(window, 'location', {
      value: { protocol: 'http:', host: 'localhost:8080' },
      configurable: true,
    });
  });

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

  function latestSocket() {
    return MockWebSocket.instances.at(-1)!;
  }

  function wrapper({ children }: { children: React.ReactNode }) {
    return <WebSocketProvider>{children}</WebSocketProvider>;
  }

  it('connects to ws dashboard with auth token query param', () => {
    setStoredAuth('drjones', 'secret');
    const { result } = renderHook(() => useWebSocketContext(), { wrapper });

    const ws = latestSocket();
    const token = btoa('drjones:secret');
    expect(ws.url).toBe(`ws://localhost:8080/ws/dashboard?token=${encodeURIComponent(token)}`);

    act(() => ws.emitOpen());
    expect(result.current.isConnected).toBe(true);
  });

  it('connects without token when logged out', () => {
    clearStoredAuth();
    renderHook(() => useWebSocketContext(), { wrapper });
    expect(latestSocket().url).toBe('ws://localhost:8080/ws/dashboard');
  });

  it('useWebSocket re-exports context hook', () => {
    expect(useWebSocket).toBe(useWebSocketContext);
  });

  it('handles init and agent_online messages', () => {
    const agent = mockAgent({ id: 'live-1' });
    const { result } = renderHook(() => useWebSocketContext(), { wrapper });

    act(() => {
      latestSocket().emitOpen();
      latestSocket().emitMessage({
        type: 'init',
        payload: { agents: [agent] },
      });
    });
    expect(result.current.agents).toEqual([agent]);

    const updated = { ...agent, hashrate_15s: 999 };
    act(() =>
      latestSocket().emitMessage({
        type: 'agent_online',
        payload: updated,
      })
    );
    expect(result.current.agents[0].hashrate_15s).toBe(999);
  });

  it('marks agent offline and caps recent shares', () => {
    const agent = mockAgent({ id: 'a-offline' });
    const { result } = renderHook(() => useWebSocketContext(), { wrapper });

    act(() => {
      latestSocket().emitOpen();
      latestSocket().emitMessage({ type: 'init', payload: { agents: [agent] } });
      latestSocket().emitMessage({
        type: 'agent_offline',
        payload: { agent_id: 'a-offline' },
      });
    });
    expect(result.current.agents[0].status).toBe('offline');

    act(() => {
      for (let i = 0; i < 55; i++) {
        latestSocket().emitMessage({
          type: 'new_share',
          payload: mockShare({ id: i }),
        });
      }
    });
    expect(result.current.recentShares.length).toBeLessThanOrEqual(50);
  });

  it('assigns monotonic _seq on command_result', () => {
    const { result } = renderHook(() => useWebSocketContext(), { wrapper });

    act(() => {
      latestSocket().emitOpen();
      latestSocket().emitMessage({
        type: 'command_result',
        payload: { agent_id: 'a1', action: 'pause', success: true },
      });
      latestSocket().emitMessage({
        type: 'command_result',
        payload: { agent_id: 'a1', action: 'get_log', success: true, message: 'log data' },
      });
    });

    expect(result.current.commandResults).toHaveLength(2);
    expect(result.current.commandResults[0]._seq).toBe(1);
    expect(result.current.commandResults[1]._seq).toBe(2);
    expect(result.current.agentLogs.a1).toBe('log data');
  });

  it('schedules reconnect after close', () => {
    vi.useFakeTimers();
    renderHook(() => useWebSocketContext(), { wrapper });
    const first = latestSocket();

    act(() => first.close());
    expect(MockWebSocket.instances).toHaveLength(1);

    act(() => vi.advanceTimersByTime(3000));
    expect(MockWebSocket.instances).toHaveLength(2);
    vi.useRealTimers();
  });

  it('closes socket on unmount', () => {
    const closeSpy = vi.spyOn(MockWebSocket.prototype, 'close');
    const { unmount } = render(<WebSocketProvider><span /></WebSocketProvider>);
    unmount();
    expect(closeSpy).toHaveBeenCalled();
  });
});