WebRTC-Broadcast/tests/integration.test.ts

import { test, expect, describe, beforeEach, afterEach } from "bun:test";
import { ClientManager } from "../src/services/ClientManager.ts";
import { SignalingService } from "../src/services/SignalingService.ts";

interface MockWebSocketMessage {
  type: string;
  data?: any;
  senderId?: string;
}

class MockWebSocket {
  private messageHandler: ((event: { data: string }) => void) | null = null;
  private openHandler: (() => void) | null = null;
  private closeHandler: (() => void) | null = null;
  
  public sentMessages: string[] = [];
  public readyState = WebSocket.OPEN;
  public url: string;

  constructor(url: string) {
    this.url = url;
    // Simulate connection opening
    setTimeout(() => {
      if (this.openHandler) {
        this.openHandler();
      }
    }, 10);
  }

  send(data: string) {
    this.sentMessages.push(data);
  }

  close() {
    this.readyState = WebSocket.CLOSED;
    if (this.closeHandler) {
      this.closeHandler();
    }
  }

  set onmessage(handler: (event: { data: string }) => void) {
    this.messageHandler = handler;
  }

  set onopen(handler: () => void) {
    this.openHandler = handler;
  }

  set onclose(handler: () => void) {
    this.closeHandler = handler;
  }

  set onerror(handler: (error: any) => void) {
    // Mock implementation
  }

  simulateMessage(message: MockWebSocketMessage) {
    if (this.messageHandler) {
      this.messageHandler({ data: JSON.stringify(message) });
    }
  }

  getLastMessage(): MockWebSocketMessage | null {
    if (this.sentMessages.length === 0) return null;
    return JSON.parse(this.sentMessages[this.sentMessages.length - 1]);
  }

  getAllMessages(): MockWebSocketMessage[] {
    return this.sentMessages.map(msg => JSON.parse(msg));
  }
}

describe("WebSocket Signaling Integration", () => {
  let clientManager: ClientManager;
  let signalingService: SignalingService;
  
  beforeEach(() => {
    clientManager = new ClientManager();
    signalingService = new SignalingService(clientManager);
  });

  test("should handle complete publisher-subscriber flow", () => {
    // Setup subscribers first
    const subscriber1Ws = new MockWebSocket("ws://test?role=subscriber");
    const subscriber1Id = signalingService.handleConnection(subscriber1Ws, "subscriber");
    
    const subscriber2Ws = new MockWebSocket("ws://test?role=subscriber");  
    const subscriber2Id = signalingService.handleConnection(subscriber2Ws, "subscriber");

    // Setup publisher (this should notify existing subscribers)
    const publisherWs = new MockWebSocket("ws://test?role=publisher");
    const publisherId = signalingService.handleConnection(publisherWs, "publisher");

    // Verify initial connections
    expect(publisherWs.getLastMessage()).toEqual({
      type: 'join',
      data: { clientId: publisherId, role: 'publisher' }
    });

    // Verify subscribers get notified about publisher joining
    const sub1Messages = subscriber1Ws.getAllMessages();
    const sub2Messages = subscriber2Ws.getAllMessages();
    
    expect(sub1Messages).toHaveLength(2); // join + publisher-joined
    expect(sub1Messages.some(msg => msg.type === 'publisher-joined')).toBe(true);
    expect(sub2Messages).toHaveLength(2); // join + publisher-joined  
    expect(sub2Messages.some(msg => msg.type === 'publisher-joined')).toBe(true);

    // Test offer flow
    const offerMessage = {
      type: "offer" as const,
      data: { sdp: "fake-offer-sdp", type: "offer" }
    };

    signalingService.handleMessage(publisherId, offerMessage);

    // Verify subscribers received the offer
    const sub1LastMessage = subscriber1Ws.getLastMessage();
    const sub2LastMessage = subscriber2Ws.getLastMessage();

    expect(sub1LastMessage).toEqual({
      ...offerMessage,
      senderId: publisherId
    });
    expect(sub2LastMessage).toEqual({
      ...offerMessage,
      senderId: publisherId
    });

    // Test answer flow
    const answerMessage = {
      type: "answer" as const,
      data: { sdp: "fake-answer-sdp", type: "answer" }
    };

    signalingService.handleMessage(subscriber1Id, answerMessage);

    // Verify publisher received the answer
    const publisherLastMessage = publisherWs.getLastMessage();
    expect(publisherLastMessage).toEqual({
      ...answerMessage,
      senderId: subscriber1Id
    });

    // Test ICE candidate exchange
    const iceCandidateMessage = {
      type: "ice-candidate" as const,
      data: { candidate: "fake-ice-candidate" }
    };

    // Publisher to subscribers
    signalingService.handleMessage(publisherId, iceCandidateMessage);
    
    expect(subscriber1Ws.getLastMessage()).toEqual({
      ...iceCandidateMessage,
      senderId: publisherId
    });
    expect(subscriber2Ws.getLastMessage()).toEqual({
      ...iceCandidateMessage,
      senderId: publisherId
    });

    // Subscriber to publisher
    signalingService.handleMessage(subscriber1Id, iceCandidateMessage);
    
    expect(publisherWs.getLastMessage()).toEqual({
      ...iceCandidateMessage,
      senderId: subscriber1Id
    });
  });

  test("should handle publisher disconnect gracefully", () => {
    // Setup 
    const publisherWs = new MockWebSocket("ws://test?role=publisher");
    const publisherId = signalingService.handleConnection(publisherWs, "publisher");
    
    const subscriberWs = new MockWebSocket("ws://test?role=subscriber");
    signalingService.handleConnection(subscriberWs, "subscriber");

    // Disconnect publisher
    signalingService.handleDisconnection(publisherId);

    // Verify subscriber gets publisher-left notification
    const subscriberMessages = subscriberWs.getAllMessages();
    expect(subscriberMessages).toHaveLength(2); // join + publisher-left
    expect(subscriberMessages.some(msg => msg.type === 'publisher-left')).toBe(true);

    // Verify publisher is removed
    expect(clientManager.getPublisher()).toBeNull();
  });

  test("should handle subscriber disconnect gracefully", () => {
    const publisherWs = new MockWebSocket("ws://test?role=publisher"); 
    signalingService.handleConnection(publisherWs, "publisher");
    
    const subscriberWs = new MockWebSocket("ws://test?role=subscriber");
    const subscriberId = signalingService.handleConnection(subscriberWs, "subscriber");

    // Verify subscriber was added
    expect(clientManager.getSubscribers()).toHaveLength(1);

    // Disconnect subscriber
    signalingService.handleDisconnection(subscriberId);

    // Verify subscriber was removed
    expect(clientManager.getSubscribers()).toHaveLength(0);
    expect(clientManager.getClient(subscriberId)).toBeUndefined();
  });

  test("should handle multiple subscribers connecting and disconnecting", () => {
    const publisherWs = new MockWebSocket("ws://test?role=publisher");
    signalingService.handleConnection(publisherWs, "publisher");

    // Add 3 subscribers
    const subscribers = [];
    for (let i = 0; i < 3; i++) {
      const ws = new MockWebSocket(`ws://test?role=subscriber&id=${i}`);
      const id = signalingService.handleConnection(ws, "subscriber");
      subscribers.push({ ws, id });
    }

    expect(clientManager.getSubscribers()).toHaveLength(3);

    // Disconnect middle subscriber
    signalingService.handleDisconnection(subscribers[1].id);
    expect(clientManager.getSubscribers()).toHaveLength(2);

    // Verify remaining subscribers are still connected
    expect(clientManager.getClient(subscribers[0].id)).toBeDefined();
    expect(clientManager.getClient(subscribers[2].id)).toBeDefined();
    expect(clientManager.getClient(subscribers[1].id)).toBeUndefined();
  });
});