WebRTC-Broadcast/tests/sfu/SFUSignalingService.test.ts

import { test, expect, describe, beforeEach } from "bun:test";
import { SFUSignalingService } from "../../src/services/SFUSignalingService.ts";
import { SFUClientManager } from "../../src/services/SFUClientManager.ts";
import type { ISFUSignalingMessage } from "../../src/interfaces/ISFUTypes.ts";

// Mock MediaServerManager
class MockMediaServerManager {
  private mockRtpCapabilities = {
    codecs: [
      { mimeType: 'video/VP8', clockRate: 90000 },
      { mimeType: 'audio/opus', clockRate: 48000 }
    ]
  };

  private mockTransport = {
    id: 'transport-123',
    iceParameters: { usernameFragment: 'test', password: 'test' },
    iceCandidates: [],
    dtlsParameters: { fingerprints: [{ algorithm: 'sha-256', value: 'test' }] },
    connect: () => Promise.resolve(),
    produce: () => Promise.resolve({ id: 'producer-123', kind: 'video' }),
    consume: () => Promise.resolve({ id: 'consumer-123', kind: 'video', rtpParameters: {} }),
    close: () => {}
  };

  private mockRouter = {
    canConsume: () => true
  };

  getRtpCapabilities() {
    return this.mockRtpCapabilities;
  }

  async createWebRtcTransport() {
    return this.mockTransport;
  }

  getRouter() {
    return this.mockRouter;
  }

  getWorker() {
    return { pid: 12345 };
  }
}

// Mock WebSocket
class MockWebSocket {
  sentMessages: string[] = [];
  
  send(data: string) {
    this.sentMessages.push(data);
  }

  getLastMessage() {
    return this.sentMessages.length > 0 ? 
      JSON.parse(this.sentMessages[this.sentMessages.length - 1]) : null;
  }

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

describe("SFUSignalingService", () => {
  let signalingService: SFUSignalingService;
  let clientManager: SFUClientManager;
  let mediaServer: MockMediaServerManager;

  beforeEach(() => {
    clientManager = new SFUClientManager();
    mediaServer = new MockMediaServerManager();
    signalingService = new SFUSignalingService(mediaServer as any, clientManager);
  });

  test("should handle client connections", () => {
    const mockWs = new MockWebSocket();
    
    const clientId = signalingService.handleConnection(mockWs, "publisher");
    
    expect(typeof clientId).toBe("string");
    expect(clientId.length).toBeGreaterThan(0);
    
    const client = clientManager.getClient(clientId);
    expect(client).toBeTruthy();
    expect(client?.role).toBe("publisher");
    expect(client?.ws).toBe(mockWs);
    
    const lastMessage = mockWs.getLastMessage();
    expect(lastMessage.type).toBe("join");
    expect(lastMessage.data.clientId).toBe(clientId);
    expect(lastMessage.data.role).toBe("publisher");
  });

  test("should handle client disconnection", () => {
    const mockWs = new MockWebSocket();
    
    const clientId = signalingService.handleConnection(mockWs, "subscriber");
    expect(clientManager.getClient(clientId)).toBeTruthy();
    
    signalingService.handleDisconnection(clientId);
    expect(clientManager.getClient(clientId)).toBeUndefined();
  });

  test("should handle getRouterRtpCapabilities message", async () => {
    const mockWs = new MockWebSocket();
    const clientId = signalingService.handleConnection(mockWs, "publisher");
    
    const message: ISFUSignalingMessage = {
      type: "getRouterRtpCapabilities"
    };
    
    await signalingService.handleMessage(clientId, message);
    
    const lastMessage = mockWs.getLastMessage();
    expect(lastMessage.type).toBe("routerRtpCapabilities");
    expect(lastMessage.data.rtpCapabilities).toBeTruthy();
    expect(lastMessage.data.rtpCapabilities.codecs).toHaveLength(2);
  });

  test("should handle createWebRtcTransport message", async () => {
    const mockWs = new MockWebSocket();
    const clientId = signalingService.handleConnection(mockWs, "publisher");
    
    const message: ISFUSignalingMessage = {
      type: "createWebRtcTransport"
    };
    
    await signalingService.handleMessage(clientId, message);
    
    const lastMessage = mockWs.getLastMessage();
    expect(lastMessage.type).toBe("webRtcTransportCreated");
    expect(lastMessage.data.id).toBe("transport-123");
    expect(lastMessage.data.iceParameters).toBeTruthy();
    expect(lastMessage.data.dtlsParameters).toBeTruthy();
    
    // Verify transport was set on client
    const client = clientManager.getClient(clientId);
    expect(client?.transport).toBeTruthy();
  });

  test("should handle connectWebRtcTransport message", async () => {
    const mockWs = new MockWebSocket();
    const clientId = signalingService.handleConnection(mockWs, "publisher");
    
    // First create transport
    await signalingService.handleMessage(clientId, { type: "createWebRtcTransport" });
    
    // Then connect it
    const connectMessage: ISFUSignalingMessage = {
      type: "connectWebRtcTransport",
      data: {
        dtlsParameters: { fingerprints: [{ algorithm: 'sha-256', value: 'test' }] }
      }
    };
    
    await signalingService.handleMessage(clientId, connectMessage);
    
    const lastMessage = mockWs.getLastMessage();
    expect(lastMessage.type).toBe("webRtcTransportConnected");
  });

  test("should handle produce message for publisher", async () => {
    const mockWs = new MockWebSocket();
    const clientId = signalingService.handleConnection(mockWs, "publisher");
    
    // Create transport first
    await signalingService.handleMessage(clientId, { type: "createWebRtcTransport" });
    
    // Mock producer creation
    const mockProducer = { id: 'producer-123', kind: 'video' };
    const client = clientManager.getClient(clientId);
    if (client?.transport) {
      (client.transport as any).produce = () => Promise.resolve(mockProducer);
    }
    
    const produceMessage: ISFUSignalingMessage = {
      type: "produce",
      data: {
        kind: "video",
        rtpParameters: { codecs: [], encodings: [] }
      }
    };
    
    await signalingService.handleMessage(clientId, produceMessage);
    
    const messages = mockWs.getAllMessages();
    const producedMessage = messages.find(msg => msg.type === 'produced');
    expect(producedMessage).toBeTruthy();
    expect(producedMessage.data.producerId).toBe('producer-123');
  });

  test("should reject produce message for subscriber", async () => {
    const mockWs = new MockWebSocket();
    const clientId = signalingService.handleConnection(mockWs, "subscriber");
    
    await signalingService.handleMessage(clientId, { type: "createWebRtcTransport" });
    
    const produceMessage: ISFUSignalingMessage = {
      type: "produce",
      data: {
        kind: "video",
        rtpParameters: { codecs: [], encodings: [] }
      }
    };
    
    await signalingService.handleMessage(clientId, produceMessage);
    
    const lastMessage = mockWs.getLastMessage();
    expect(lastMessage.type).toBe("error");
    expect(lastMessage.data.message).toContain("Only publishers can produce");
  });

  test("should handle consume message for subscriber", async () => {
    // First create a publisher with a producer
    const publisherWs = new MockWebSocket();
    const publisherId = signalingService.handleConnection(publisherWs, "publisher");
    await signalingService.handleMessage(publisherId, { type: "createWebRtcTransport" });
    
    // Mock producer
    const mockProducer = { id: 'producer-123', kind: 'video' };
    clientManager.setClientProducer(publisherId, mockProducer as any);
    
    // Now create subscriber
    const subscriberWs = new MockWebSocket();
    const subscriberId = signalingService.handleConnection(subscriberWs, "subscriber");
    await signalingService.handleMessage(subscriberId, { type: "createWebRtcTransport" });
    
    // Mock consumer creation
    const mockConsumer = { 
      id: 'consumer-123', 
      kind: 'video', 
      rtpParameters: { codecs: [], encodings: [] }
    };
    const subscriberClient = clientManager.getClient(subscriberId);
    if (subscriberClient?.transport) {
      (subscriberClient.transport as any).consume = () => Promise.resolve(mockConsumer);
    }
    
    const consumeMessage: ISFUSignalingMessage = {
      type: "consume",
      data: {
        producerId: "producer-123",
        rtpCapabilities: { codecs: [] }
      }
    };
    
    await signalingService.handleMessage(subscriberId, consumeMessage);
    
    const lastMessage = subscriberWs.getLastMessage();
    expect(lastMessage.type).toBe("consumed");
    expect(lastMessage.data.consumerId).toBe('consumer-123');
    expect(lastMessage.data.producerId).toBe('producer-123');
  });

  test("should reject consume message for publisher", async () => {
    const mockWs = new MockWebSocket();
    const clientId = signalingService.handleConnection(mockWs, "publisher");
    
    await signalingService.handleMessage(clientId, { type: "createWebRtcTransport" });
    
    const consumeMessage: ISFUSignalingMessage = {
      type: "consume",
      data: {
        producerId: "producer-123",
        rtpCapabilities: { codecs: [] }
      }
    };
    
    await signalingService.handleMessage(clientId, consumeMessage);
    
    const lastMessage = mockWs.getLastMessage();
    expect(lastMessage.type).toBe("error");
    expect(lastMessage.data.message).toContain("Only subscribers can consume");
  });

  test("should handle resume message", async () => {
    const mockWs = new MockWebSocket();
    const clientId = signalingService.handleConnection(mockWs, "subscriber");
    
    // Mock consumer
    const mockConsumer = { 
      id: 'consumer-123',
      resume: () => Promise.resolve()
    };
    clientManager.addClientConsumer(clientId, mockConsumer as any);
    
    const resumeMessage: ISFUSignalingMessage = {
      type: "resume",
      data: { consumerId: "consumer-123" }
    };
    
    await signalingService.handleMessage(clientId, resumeMessage);
    
    const lastMessage = mockWs.getLastMessage();
    expect(lastMessage.type).toBe("resumed");
    expect(lastMessage.data.consumerId).toBe("consumer-123");
  });

  test("should handle pause message", async () => {
    const mockWs = new MockWebSocket();
    const clientId = signalingService.handleConnection(mockWs, "subscriber");
    
    // Mock consumer
    const mockConsumer = { 
      id: 'consumer-123',
      pause: () => Promise.resolve()
    };
    clientManager.addClientConsumer(clientId, mockConsumer as any);
    
    const pauseMessage: ISFUSignalingMessage = {
      type: "pause",
      data: { consumerId: "consumer-123" }
    };
    
    await signalingService.handleMessage(clientId, pauseMessage);
    
    const lastMessage = mockWs.getLastMessage();
    expect(lastMessage.type).toBe("paused");
    expect(lastMessage.data.consumerId).toBe("consumer-123");
  });

  test("should handle getProducers message", async () => {
    // Create publisher with producer
    const publisherWs = new MockWebSocket();
    const publisherId = signalingService.handleConnection(publisherWs, "publisher");
    
    const mockProducer = { id: 'producer-123', kind: 'video' };
    clientManager.setClientProducer(publisherId, mockProducer as any);
    
    // Create subscriber requesting producers
    const subscriberWs = new MockWebSocket();
    const subscriberId = signalingService.handleConnection(subscriberWs, "subscriber");
    
    const getProducersMessage: ISFUSignalingMessage = {
      type: "getProducers"
    };
    
    await signalingService.handleMessage(subscriberId, getProducersMessage);
    
    const lastMessage = subscriberWs.getLastMessage();
    expect(lastMessage.type).toBe("producers");
    expect(lastMessage.data.producers).toHaveLength(1);
    expect(lastMessage.data.producers[0].id).toBe('producer-123');
    expect(lastMessage.data.producers[0].kind).toBe('video');
  });

  test("should notify subscribers when new producer joins", async () => {
    // Create subscribers first
    const subscriber1Ws = new MockWebSocket();
    const subscriber1Id = signalingService.handleConnection(subscriber1Ws, "subscriber");
    
    const subscriber2Ws = new MockWebSocket();
    const subscriber2Id = signalingService.handleConnection(subscriber2Ws, "subscriber");
    
    // Clear initial join messages
    subscriber1Ws.sentMessages = [];
    subscriber2Ws.sentMessages = [];
    
    // Create publisher and produce
    const publisherWs = new MockWebSocket();
    const publisherId = signalingService.handleConnection(publisherWs, "publisher");
    await signalingService.handleMessage(publisherId, { type: "createWebRtcTransport" });
    
    // Mock producer creation and produce
    const mockProducer = { id: 'producer-123', kind: 'video' };
    const publisherClient = clientManager.getClient(publisherId);
    if (publisherClient?.transport) {
      (publisherClient.transport as any).produce = () => Promise.resolve(mockProducer);
    }
    
    await signalingService.handleMessage(publisherId, {
      type: "produce",
      data: { kind: "video", rtpParameters: { codecs: [], encodings: [] } }
    });
    
    // Check that subscribers were notified
    const sub1LastMessage = subscriber1Ws.getLastMessage();
    const sub2LastMessage = subscriber2Ws.getLastMessage();
    
    expect(sub1LastMessage?.type).toBe("newProducer");
    expect(sub1LastMessage?.data.producerId).toBe('producer-123');
    
    expect(sub2LastMessage?.type).toBe("newProducer");
    expect(sub2LastMessage?.data.producerId).toBe('producer-123');
  });

  test("should handle messages from unknown clients gracefully", async () => {
    const message: ISFUSignalingMessage = {
      type: "getRouterRtpCapabilities"
    };
    
    expect(async () => {
      await signalingService.handleMessage("unknown-client-id", message);
    }).not.toThrow();
  });

  test("should handle unknown message types", async () => {
    const mockWs = new MockWebSocket();
    const clientId = signalingService.handleConnection(mockWs, "publisher");
    
    const unknownMessage: any = {
      type: "unknownMessageType"
    };
    
    await signalingService.handleMessage(clientId, unknownMessage);
    
    // Should not crash, just log a warning
    expect(true).toBe(true); // Test passes if no exception thrown
  });

  test("should provide service statistics", () => {
    // Create some clients
    const publisher1Ws = new MockWebSocket();
    const publisher1Id = signalingService.handleConnection(publisher1Ws, "publisher");
    
    const subscriber1Ws = new MockWebSocket();
    const subscriber1Id = signalingService.handleConnection(subscriber1Ws, "subscriber");
    
    const subscriber2Ws = new MockWebSocket();
    const subscriber2Id = signalingService.handleConnection(subscriber2Ws, "subscriber");
    
    // Add some producers and consumers
    const mockProducer = { id: 'producer-123', kind: 'video' };
    const mockConsumer1 = { id: 'consumer-123', kind: 'video' };
    const mockConsumer2 = { id: 'consumer-456', kind: 'video' };
    
    clientManager.setClientProducer(publisher1Id, mockProducer as any);
    clientManager.addClientConsumer(subscriber1Id, mockConsumer1 as any);
    clientManager.addClientConsumer(subscriber2Id, mockConsumer2 as any);
    
    const stats = signalingService.getStats();
    
    expect(stats.clients).toBe(3);
    expect(stats.publishers).toBe(1);
    expect(stats.subscribers).toBe(2);
    expect(stats.producers).toBe(1);
    expect(stats.consumers).toBe(2);
    expect(stats.mediaServer.workerId).toBe(12345);
  });
});