welshman/packages/net/__tests__/Socket.test.ts

import { sleep } from "@welshman/lib"
import WebSocket from 'isomorphic-ws'
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
import { Socket, SocketStatus, SocketEventType } from "../src/socket"
import { ClientMessage, RelayMessage } from "../src/message"

vi.mock('isomorphic-ws', () => {
  const WebSocket = vi.fn(function () {
    setTimeout(() => this.onopen())
  })

  WebSocket.prototype.send = vi.fn()

  WebSocket.prototype.close = vi.fn(function () {
    this.onclose()
  })

  return { default: WebSocket }
})

describe("Socket", () => {
  let socket: Socket

  beforeEach(() => {
    vi.useFakeTimers()
    socket = new Socket("wss://test.relay")
  })

  afterEach(() => {
    vi.clearAllMocks()
    vi.useRealTimers()
    socket.cleanup()
  })

  it("should initialize with correct url", () => {
    expect(socket.url).toBe("wss://test.relay")
  })

  describe("open", () => {
    it("should create websocket and emit opening status", () => {
      const statusSpy = vi.fn()
      socket.on(SocketEventType.Status, statusSpy)

      socket.open()

      expect(socket._ws).toBeDefined()
      expect(statusSpy).toHaveBeenCalledWith(SocketStatus.Opening, "wss://test.relay")

      vi.runAllTimers()

      expect(statusSpy).toHaveBeenCalledWith(SocketStatus.Open, "wss://test.relay")
    })

    it("should throw error if socket already exists", () => {
      socket.open()
      expect(() => socket.open()).toThrow("Attempted to open a websocket that has not been closed")
    })

    it("should emit invalid status on invalid URL", () => {
      const statusSpy = vi.fn()
      socket.on(SocketEventType.Status, statusSpy)

      vi.mocked(WebSocket).mockImplementationOnce(() => {
        throw new Error()
      })

      socket.open()

      expect(statusSpy).toHaveBeenCalledWith(SocketStatus.Invalid, "wss://test.relay")
    })
  })

  describe("close", () => {
    it("should close websocket and emit closed status", () => {
      const statusSpy = vi.fn()
      socket.on(SocketEventType.Status, statusSpy)

      socket.open()

      const ws = socket._ws

      socket.close()

      expect(ws.close).toHaveBeenCalled()
      expect(statusSpy).toHaveBeenCalledWith(SocketStatus.Closed, "wss://test.relay")
    })
  })

  describe("send", () => {
    it("should queue messages and emit enqueue event", () => {
      const enqueueSpy = vi.fn()
      socket.on(SocketEventType.Enqueue, enqueueSpy)

      const message: ClientMessage = ["EVENT", { id: "123", kind: 1 }]
      socket.send(message)

      expect(enqueueSpy).toHaveBeenCalledWith(message, "wss://test.relay")
    })

    it("should send messages when socket is open", async () => {
      const sendSpy = vi.fn()
      socket.on(SocketEventType.Send, sendSpy)

      socket.open()
      socket._ws.onopen()

      const message: ClientMessage = ["EVENT", { id: "123", kind: 1 }]
      socket.send(message)

      await vi.runAllTimers()

      expect(socket._ws.send).toHaveBeenCalledWith(JSON.stringify(message))
      expect(sendSpy).toHaveBeenCalledWith(message, "wss://test.relay")
    })
  })

  describe("receive", () => {
    it("should handle valid relay messages", async () => {
      const receiveSpy = vi.fn()
      socket.on(SocketEventType.Receive, receiveSpy)

      socket.open()
      const message: RelayMessage = ["EVENT", "123", { id: "123", kind: 1 }]
      socket._ws.onmessage({ data: JSON.stringify(message) })

      // Allow task queue to process
      await vi.runAllTimers()

      expect(receiveSpy).toHaveBeenCalledWith(message, "wss://test.relay")
    })

    it("should emit error on invalid JSON", () => {
      const errorSpy = vi.fn()
      socket.on(SocketEventType.Error, errorSpy)

      socket.open()
      socket._ws.onmessage({ data: "invalid json" })

      expect(errorSpy).toHaveBeenCalledWith("Invalid message received", "wss://test.relay")
    })

    it("should emit error on non-array message", () => {
      const errorSpy = vi.fn()
      socket.on(SocketEventType.Error, errorSpy)

      socket.open()
      socket._ws.onmessage({ data: JSON.stringify({ not: "an array" }) })

      expect(errorSpy).toHaveBeenCalledWith("Invalid message received", "wss://test.relay")
    })
  })

  describe("cleanup", () => {
    it("should close socket and clear queues", () => {
      socket.open()

      const ws = socket._ws

      socket.cleanup()

      expect(ws.close).toHaveBeenCalled()
      expect(socket.listenerCount(SocketEventType.Send)).toBe(0)
    })
  })

  describe("error handling", () => {
    it("should emit error status on websocket error", () => {
      const statusSpy = vi.fn()
      socket.on(SocketEventType.Status, statusSpy)

      socket.open()
      socket._ws.onerror()

      expect(statusSpy).toHaveBeenCalledWith(SocketStatus.Error, "wss://test.relay")
    })
  })

  describe("attemptToOpen", () => {
    it("should open socket if not already open", () => {
      const openSpy = vi.spyOn(socket, "open")

      socket.attemptToOpen()
      expect(openSpy).toHaveBeenCalled()
    })

    it("should not open socket if already open", () => {
      const openSpy = vi.spyOn(socket, "open")

      socket.open()
      socket.attemptToOpen()

      expect(openSpy).toHaveBeenCalledTimes(1)
    })
  })
})