workavia-calendar-front/__test__/features/websocket/createWebSocketConnection.test.tsx

import { fetchWebSocketTicket } from '@/websocket/api/fetchWebSocketTicket'
import { createWebSocketConnection } from '@/websocket/connection/createConnection'
import { WS_INBOUND_EVENTS } from '@/websocket/protocols'
import { waitFor } from '@testing-library/dom'
import { setupWebsocket } from './utils/setupWebsocket'

jest.mock('@/websocket/api/fetchWebSocketTicket')

describe('createWebSocketConnection', () => {
  let mockWebSocket: jest.Mock
  let cleanup: () => void
  let webSocketInstances: any[] = []

  const mockTicket = {
    value: 'test-ticket-123',
    clientAddress: '127.0.0.1',
    generatedOn: '2025-01-12T10:00:00Z',
    validUntil: '2025-01-12T11:00:00Z',
    username: 'testuser'
  }

  /** ---------- Helpers ---------- */

  const getWs = () => webSocketInstances[0]

  const triggerEvent = (ws: any, event: string, payload?: any) => {
    ws._listeners[event]?.[0]?.(payload)
  }

  const createAndOpenConnection = async () => {
    const mockCallbacks = {
      onMessage: jest.fn(),
      onClose: jest.fn(),
      onError: jest.fn()
    }
    const promise = createWebSocketConnection(mockCallbacks)

    await waitFor(() => {
      expect(webSocketInstances.length).toBe(1)
    })

    triggerEvent(getWs(), WS_INBOUND_EVENTS.CONNECTION_OPENED)
    const socket = await promise

    return { socket, ws: getWs(), promise, mockCallbacks }
  }

  /** ---------- Setup ---------- */

  beforeEach(() => {
    ;({ webSocketInstances, mockWebSocket, cleanup } = setupWebsocket())
    window.WEBSOCKET_URL = 'wss://calendar.example.com'
    ;(fetchWebSocketTicket as jest.Mock).mockResolvedValue(mockTicket)
  })

  afterEach(() => {
    jest.clearAllMocks()
    delete window.WEBSOCKET_URL
    delete window.CALENDAR_BASE_URL
    cleanup()
  })

  /** ---------- Tests ---------- */

  it('throws when WEBSOCKET_URL is not defined', async () => {
    delete window.WEBSOCKET_URL
    const mockCallbacks = {
      onMessage: jest.fn()
    }

    await expect(createWebSocketConnection(mockCallbacks)).rejects.toThrow(
      'WEBSOCKET_URL is not defined'
    )
  })

  it('fetches WebSocket ticket', async () => {
    await createAndOpenConnection()
    expect(fetchWebSocketTicket).toHaveBeenCalledTimes(1)
  })

  it('creates WebSocket with correct URL and ticket', async () => {
    await createAndOpenConnection()

    expect(mockWebSocket).toHaveBeenCalledWith(
      'wss://calendar.example.com/ws?ticket=test-ticket-123'
    )
  })

  it('creates WebSocket with correct URL and ticket without the WEBSOCKET_URL', async () => {
    delete window.WEBSOCKET_URL
    window.CALENDAR_BASE_URL = 'https://calendar.example.com'
    await createAndOpenConnection()

    expect(mockWebSocket).toHaveBeenCalledWith(
      'wss://calendar.example.com/ws?ticket=test-ticket-123'
    )
  })

  it('resolves with socket when connection opens', async () => {
    const { socket, ws } = await createAndOpenConnection()
    expect(socket).toBe(ws)
  })

  it('rejects when connection fails', async () => {
    const mockCallbacks = {
      onMessage: jest.fn()
    }
    const promise = createWebSocketConnection(mockCallbacks)

    await waitFor(() => {
      expect(webSocketInstances.length).toBe(1)
    })

    triggerEvent(getWs(), WS_INBOUND_EVENTS.ERROR, new Event('error'))

    await expect(promise).rejects.toThrow('WebSocket connection failed')
  })

  it('attaches message event listener', async () => {
    const { ws } = await createAndOpenConnection()

    expect(ws._listeners[WS_INBOUND_EVENTS.MESSAGE]).toBeDefined()
    expect(ws._listeners[WS_INBOUND_EVENTS.MESSAGE].length).toBeGreaterThan(0)
  })

  it('attaches close event listener', async () => {
    const { ws } = await createAndOpenConnection()
    expect(ws._listeners[WS_INBOUND_EVENTS.CONNECTION_CLOSED]).toBeDefined()
  })

  it('handles invalid JSON messages', async () => {
    const errorSpy = jest.spyOn(console, 'error').mockImplementation()

    const { ws } = await createAndOpenConnection()

    triggerEvent(ws, WS_INBOUND_EVENTS.MESSAGE, { data: 'invalid json' })

    expect(errorSpy).toHaveBeenCalledWith(
      'Failed to parse WebSocket message:',
      expect.any(Error)
    )

    errorSpy.mockRestore()
  })

  it('rejects on timeout', async () => {
    jest.useFakeTimers()
    const mockCallbacks = {
      onMessage: jest.fn()
    }
    const promise = createWebSocketConnection(mockCallbacks)

    await waitFor(() => {
      expect(webSocketInstances.length).toBe(1)
    })

    jest.advanceTimersByTime(10000)

    await expect(promise).rejects.toThrow('WebSocket connection timed out')

    jest.useRealTimers()
  })

  it('calls onMessage callback when message received', async () => {
    const { ws, mockCallbacks } = await createAndOpenConnection()

    const testMessage = { type: 'test', payload: 'data' }
    triggerEvent(ws, WS_INBOUND_EVENTS.MESSAGE, {
      data: JSON.stringify(testMessage)
    })

    expect(mockCallbacks.onMessage).toHaveBeenCalledWith(testMessage)
  })

  it('does not call onMessage when JSON parsing fails', async () => {
    const errorSpy = jest.spyOn(console, 'error').mockImplementation()
    const { ws, mockCallbacks } = await createAndOpenConnection()

    triggerEvent(ws, WS_INBOUND_EVENTS.MESSAGE, { data: 'invalid json' })

    expect(mockCallbacks.onMessage).not.toHaveBeenCalled()
    expect(errorSpy).toHaveBeenCalledWith(
      'Failed to parse WebSocket message:',
      expect.any(Error)
    )

    errorSpy.mockRestore()
  })

  it('calls onClose callback when connection closes', async () => {
    const { ws, mockCallbacks } = await createAndOpenConnection()

    const closeEvent = new CloseEvent('close', {
      code: 1000,
      reason: 'Normal closure'
    })
    triggerEvent(ws, WS_INBOUND_EVENTS.CONNECTION_CLOSED, closeEvent)

    expect(mockCallbacks.onClose).toHaveBeenCalledWith(closeEvent)
  })

  it('calls onError callback when error occurs', async () => {
    const { ws, mockCallbacks } = await createAndOpenConnection()

    const errorEvent = new Event('error')
    triggerEvent(ws, WS_INBOUND_EVENTS.ERROR, errorEvent)

    expect(mockCallbacks.onError).toHaveBeenCalledWith(errorEvent)
  })
})