workavia-calendar-front/__test__/features/websocket/utils/setupWebsocket.ts

export function setupWebsocket() {
  const originalWebSocket = global.WebSocket;
  const webSocketInstances: any[] = [];
  const mockWebSocket = jest.fn().mockImplementation((url: string) => {
    const ws = {
      url,
      readyState: WebSocket.CONNECTING,
      addEventListener: jest.fn(),
      removeEventListener: jest.fn(),
      send: jest.fn(),
      close: jest.fn(),
      _listeners: {} as Record<string, Function[]>,
    };

    ws.addEventListener.mockImplementation((event, handler) => {
      ws._listeners[event] ??= [];
      ws._listeners[event].push(handler);
    });

    ws.removeEventListener.mockImplementation((event, handler) => {
      ws._listeners[event] =
        ws._listeners[event]?.filter((h) => h !== handler) ?? [];
    });

    webSocketInstances.push(ws);
    return ws;
  });

  global.WebSocket = mockWebSocket as any;
  const cleanup = () => {
    global.WebSocket = originalWebSocket;
    webSocketInstances.length = 0;
  };
  return { webSocketInstances, mockWebSocket, cleanup };
}