useWebSocket.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { act, renderHook } from '@testing-library/react'
import { FakeWebSocket } from '../test/fakeWebSocket'
import { useWebSocket } from './useWebSocket'
import type { WsMessage } from './protocol'

describe('useWebSocket', () => {
  beforeEach(() => {
    vi.useFakeTimers()
    FakeWebSocket.reset()
    vi.stubGlobal('WebSocket', FakeWebSocket)
  })

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

  it('sends the JOIN handshake as soon as the socket opens', () => {
    renderHook(() => useWebSocket('display', () => {}))

    act(() => FakeWebSocket.latest.open())

    const join = FakeWebSocket.latest.sentMessages()[0]
    expect(join.type).toBe('JOIN')
    expect(join.sender).toBeNull()
    expect(join.payload).toEqual({ role: 'display' })
  })

  it('learns its identity from WELCOME and renames from the roster', () => {
    const { result } = renderHook(() => useWebSocket('display', () => {}))
    act(() => FakeWebSocket.latest.open())
    expect(result.current.self).toBeNull()

    act(() =>
      FakeWebSocket.latest.receive({
        type: 'WELCOME',
        sender: null,
        payload: { id: 'abc123', name: 'display-abc1', role: 'display' },
        timestamp: 1,
      }),
    )
    expect(result.current.self).toEqual({ id: 'abc123', name: 'display-abc1', role: 'display' })

    act(() =>
      FakeWebSocket.latest.receive({
        type: 'PRESENCE',
        sender: null,
        payload: { roster: [{ id: 'abc123', name: 'Nick', role: 'display' }] },
        timestamp: 2,
      }),
    )
    expect(result.current.self?.name).toBe('Nick')
    expect(result.current.roster).toHaveLength(1)
  })

  it('rename sends a RENAME frame', () => {
    const { result } = renderHook(() => useWebSocket('display', () => {}))
    act(() => FakeWebSocket.latest.open())

    result.current.rename('Sarah')

    const frame = FakeWebSocket.latest.sentMessages().at(-1)!
    expect(frame.type).toBe('RENAME')
    expect(frame.payload).toEqual({ name: 'Sarah' })
  })

  it('sendBroadcast wraps the payload in a BROADCAST envelope', () => {
    const { result } = renderHook(() => useWebSocket('display', () => {}))
    act(() => FakeWebSocket.latest.open())

    result.current.sendBroadcast({ kind: 'announcement', text: 'to all' })

    const frame = FakeWebSocket.latest.sentMessages().at(-1)!
    expect(frame.type).toBe('BROADCAST')
    expect(frame.payload).toEqual({ kind: 'announcement', text: 'to all' })
  })

  it('reports status transitions: connecting, open, closed', () => {
    const { result } = renderHook(() => useWebSocket('control', () => {}))
    expect(result.current.status).toBe('connecting')

    act(() => FakeWebSocket.latest.open())
    expect(result.current.status).toBe('open')

    act(() => FakeWebSocket.latest.drop())
    expect(result.current.status).toBe('closed')
  })

  it('delivers parsed inbound messages to the callback', () => {
    const received: WsMessage[] = []
    renderHook(() => useWebSocket('control', (msg) => received.push(msg)))
    act(() => FakeWebSocket.latest.open())

    const frame: WsMessage = {
      type: 'EVENT',
      sender: { id: 'd1', name: 'display-d1', role: 'display' },
      payload: { kind: 'reaction', emoji: '🎉' },
      timestamp: 123,
    }
    act(() => FakeWebSocket.latest.receive(frame))

    expect(received).toEqual([frame])
  })

  it('sendEvent wraps the payload in an EVENT envelope', () => {
    const { result } = renderHook(() => useWebSocket('control', () => {}))
    act(() => FakeWebSocket.latest.open())

    result.current.sendEvent({ kind: 'announcement', text: 'hi' })

    const event = FakeWebSocket.latest.sentMessages().at(-1)!
    expect(event.type).toBe('EVENT')
    expect(event.sender).toBeNull()
    expect(event.payload).toEqual({ kind: 'announcement', text: 'hi' })
  })

  it('sendEvent is a no-op while the socket is not open', () => {
    const { result } = renderHook(() => useWebSocket('control', () => {}))

    result.current.sendEvent({ kind: 'announcement', text: 'lost' })

    expect(FakeWebSocket.latest.sent).toHaveLength(0)
  })

  it('reconnects after a drop, with growing backoff', () => {
    renderHook(() => useWebSocket('display', () => {}))
    act(() => FakeWebSocket.latest.open())
    expect(FakeWebSocket.instances).toHaveLength(1)

    // First drop: reconnect after 500ms.
    act(() => FakeWebSocket.latest.drop())
    act(() => vi.advanceTimersByTime(499))
    expect(FakeWebSocket.instances).toHaveLength(1)
    act(() => vi.advanceTimersByTime(1))
    expect(FakeWebSocket.instances).toHaveLength(2)

    // Second consecutive failure (never opened): backoff doubles to 1000ms.
    act(() => FakeWebSocket.latest.drop())
    act(() => vi.advanceTimersByTime(999))
    expect(FakeWebSocket.instances).toHaveLength(2)
    act(() => vi.advanceTimersByTime(1))
    expect(FakeWebSocket.instances).toHaveLength(3)
  })

  it('a successful open resets the backoff', () => {
    renderHook(() => useWebSocket('display', () => {}))
    act(() => FakeWebSocket.latest.open())
    act(() => FakeWebSocket.latest.drop())
    act(() => vi.advanceTimersByTime(500))
    expect(FakeWebSocket.instances).toHaveLength(2)

    // Open then drop again: delay is back to the initial 500ms, not 1000ms.
    act(() => FakeWebSocket.latest.open())
    act(() => FakeWebSocket.latest.drop())
    act(() => vi.advanceTimersByTime(500))
    expect(FakeWebSocket.instances).toHaveLength(3)
  })

  it('does not reconnect after unmount', () => {
    const { unmount } = renderHook(() => useWebSocket('display', () => {}))
    act(() => FakeWebSocket.latest.open())

    unmount()
    act(() => vi.advanceTimersByTime(10_000))

    expect(FakeWebSocket.instances).toHaveLength(1)
  })
})