learning_ai_common_plat/packages/fastify-sse/src/hub.test.ts

import { describe, it, expect, beforeEach, vi } from 'vitest';
import { SSEHub } from './hub.js';
import { EventEmitter } from 'node:events';
import type { ServerResponse } from 'node:http';

function mockResponse(): ServerResponse {
  const emitter = new EventEmitter();
  const chunks: string[] = [];
  const res = Object.assign(emitter, {
    writeHead: vi.fn(),
    write: vi.fn((chunk: string) => {
      chunks.push(chunk);
      return true;
    }),
    end: vi.fn(),
    _chunks: chunks,
  });
  return res as unknown as ServerResponse;
}

describe('SSEHub', () => {
  let hub: SSEHub;

  beforeEach(() => {
    hub = new SSEHub();
  });

  describe('addClient', () => {
    it('returns a client ID and sets SSE headers', () => {
      const res = mockResponse();
      const id = hub.addClient(res);
      expect(id).toMatch(/^sse_/);
      expect(res.writeHead).toHaveBeenCalledWith(
        200,
        expect.objectContaining({
          'Content-Type': 'text/event-stream',
        })
      );
      expect(hub.clientCount).toBe(1);
    });

    it('sends initial connected event', () => {
      const res = mockResponse();
      hub.addClient(res);
      expect(res.write).toHaveBeenCalledWith(expect.stringContaining('event: connected'));
    });

    it('removes client on close', () => {
      const res = mockResponse();
      hub.addClient(res);
      expect(hub.clientCount).toBe(1);
      res.emit('close');
      expect(hub.clientCount).toBe(0);
    });
  });

  describe('broadcast', () => {
    it('sends to all connected clients', () => {
      const res1 = mockResponse();
      const res2 = mockResponse();
      hub.addClient(res1);
      hub.addClient(res2);

      const sent = hub.broadcast({ event: 'test', data: '{"hello":true}' });
      expect(sent).toBe(2);
      // Each res gets: initial connected write + broadcast write = 2 calls
      expect(res1.write).toHaveBeenCalledTimes(2);
      expect(res2.write).toHaveBeenCalledTimes(2);
    });

    it('returns 0 when no clients', () => {
      const sent = hub.broadcast({ data: 'test' });
      expect(sent).toBe(0);
    });

    it('removes clients that throw on write', () => {
      const res1 = mockResponse();
      const res2 = mockResponse();
      hub.addClient(res1);
      hub.addClient(res2);

      // Make res1 throw on next write
      (res1.write as ReturnType<typeof vi.fn>).mockImplementationOnce(() => {
        throw new Error('broken');
      });

      hub.broadcast({ data: 'test' });
      // res1 should be removed after the error
      expect(hub.clientCount).toBe(1);
    });
  });

  describe('sendToUser', () => {
    it('sends only to matching userId', () => {
      const res1 = mockResponse();
      const res2 = mockResponse();
      hub.addClient(res1, 'user-a');
      hub.addClient(res2, 'user-b');

      const sent = hub.sendToUser('user-a', { data: 'targeted' });
      expect(sent).toBe(1);
      // res1: connected + targeted = 2 writes
      expect(res1.write).toHaveBeenCalledTimes(2);
      // res2: only connected = 1 write
      expect(res2.write).toHaveBeenCalledTimes(1);
    });

    it('returns 0 when no matching users', () => {
      const res = mockResponse();
      hub.addClient(res, 'user-a');
      const sent = hub.sendToUser('user-z', { data: 'nothing' });
      expect(sent).toBe(0);
    });
  });

  describe('heartbeat', () => {
    it('sends comment to all clients', () => {
      const res = mockResponse();
      hub.addClient(res);
      hub.heartbeat();
      expect(res.write).toHaveBeenCalledWith(': heartbeat\n\n');
    });
  });

  describe('disconnectAll', () => {
    it('ends all client responses and clears', () => {
      const res1 = mockResponse();
      const res2 = mockResponse();
      hub.addClient(res1);
      hub.addClient(res2);
      expect(hub.clientCount).toBe(2);

      hub.disconnectAll();
      expect(hub.clientCount).toBe(0);
      expect(res1.end).toHaveBeenCalled();
      expect(res2.end).toHaveBeenCalled();
    });
  });

  describe('formatSSE', () => {
    it('formats with event, id, and data', () => {
      const res = mockResponse();
      hub.addClient(res);
      hub.broadcast({ event: 'task.created', data: '{}', id: 'evt_1' });

      const lastWrite = (res.write as ReturnType<typeof vi.fn>).mock.calls.at(-1)?.[0] as string;
      expect(lastWrite).toContain('id: evt_1');
      expect(lastWrite).toContain('event: task.created');
      expect(lastWrite).toContain('data: {}');
      expect(lastWrite).toMatch(/\n\n$/);
    });

    it('formats with retry field', () => {
      const res = mockResponse();
      hub.addClient(res);
      hub.broadcast({ data: 'test', retry: 5000 });

      const lastWrite = (res.write as ReturnType<typeof vi.fn>).mock.calls.at(-1)?.[0] as string;
      expect(lastWrite).toContain('retry: 5000');
    });
  });
});