WebRTC-Broadcast/tests/websocket.test.ts

import { test, expect, describe, beforeAll, afterAll } from "bun:test";
import type { ServerWebSocket } from 'bun';
import { ClientManager } from '../src/services/ClientManager.ts';
import { SignalingService } from '../src/services/SignalingService.ts';

let server: any;
let clientManager: ClientManager;
let signalingService: SignalingService;
const clientSessions = new Map<object, string>();

beforeAll(async () => {
  clientManager = new ClientManager();
  signalingService = new SignalingService(clientManager);
  
  server = Bun.serve({
    port: 0, // Use random available port
    fetch(req, server) {
      // Handle WebSocket upgrade
      const url = new URL(req.url);
      const role = url.searchParams.get('role');
      
      // Reject invalid roles before upgrade
      if (role === 'invalid') {
        return new Response('Invalid role', { status: 400 });
      }
      
      if (server.upgrade(req, { data: { role } })) {
        return; // do not return a Response
      }
      return new Response('Test Server');
    },
    websocket: {
      open(ws: ServerWebSocket<{ role?: string }>) {
        const role = ws.data?.role;
        
        if (role === 'invalid') {
          ws.close(1002, 'Invalid role parameter');
          return;
        }
        
        const clientRole = (role === 'publisher' || role === 'subscriber') ? role : 'unknown';
        const clientId = signalingService.handleConnection(ws, clientRole);
        clientSessions.set(ws, clientId);
        
        // Send join message immediately for valid roles
        if (clientRole !== 'unknown') {
          ws.send(JSON.stringify({
            type: 'join',
            data: { clientId, role: clientRole }
          }));
        }
      },
      message(ws: ServerWebSocket<unknown>, message: string | Buffer) {
        const clientId = clientSessions.get(ws);
        if (!clientId) return;
        // Message handling logic would go here
      },
      close(ws: ServerWebSocket<unknown>) {
        const clientId = clientSessions.get(ws);
        if (clientId) {
          signalingService.handleDisconnection(clientId);
          clientSessions.delete(ws);
        }
      }
    }
  });
  
  // Wait a bit for server to start
  await new Promise(resolve => setTimeout(resolve, 100));
});

afterAll(() => {
  if (server) {
    server.stop();
  }
});

describe("WebSocket Connection Test", () => {
  test("should connect to WebSocket server", async () => {
    // This test verifies the WebSocket endpoint is working
    const ws = new WebSocket(`ws://localhost:${server.port}?role=publisher`);
    
    const connectionPromise = new Promise((resolve, reject) => {
      ws.onopen = () => resolve("connected");
      ws.onerror = (error) => reject(error);
      
      // Timeout after 5 seconds
      setTimeout(() => reject(new Error("Connection timeout")), 5000);
    });

    const result = await connectionPromise;
    expect(result).toBe("connected");
    
    // Test sending a message
    const messagePromise = new Promise((resolve) => {
      ws.onmessage = (event) => {
        const message = JSON.parse(event.data);
        resolve(message);
      };
    });

    // Should receive a join message
    const joinMessage = await messagePromise;
    expect(joinMessage.type).toBe("join");
    expect(joinMessage.data.role).toBe("publisher");
    
    ws.close();
  });

  test("should handle subscriber connection", async () => {
    const ws = new WebSocket(`ws://localhost:${server.port}?role=subscriber`);
    
    const connectionPromise = new Promise((resolve, reject) => {
      ws.onopen = () => resolve("connected");
      ws.onerror = (error) => reject(error);
      setTimeout(() => reject(new Error("Connection timeout")), 5000);
    });

    await connectionPromise;
    
    const messagePromise = new Promise((resolve) => {
      ws.onmessage = (event) => {
        const message = JSON.parse(event.data);
        resolve(message);
      };
    });

    const joinMessage = await messagePromise;
    expect(joinMessage.type).toBe("join");
    expect(joinMessage.data.role).toBe("subscriber");
    
    ws.close();
  });

  test("should reject invalid role", async () => {
    const ws = new WebSocket(`ws://localhost:${server.port}?role=invalid`);
    
    const connectionPromise = new Promise((resolve, reject) => {
      ws.onopen = () => reject(new Error("Should not connect"));
      ws.onerror = (error) => {
        // WebSocket connection should fail due to 400 response
        resolve("rejected");
      };
      ws.onclose = (event) => {
        resolve("rejected");
      };
      setTimeout(() => reject(new Error("Connection timeout")), 5000);
    });

    const result = await connectionPromise;
    expect(result).toBe("rejected");
  });
});