learning_ai_common_plat/packages/fastify-sse/src/per-request.test.ts

import { describe, it, expect, vi } from 'vitest';
import { startSSE, sendSSEData, sendSSEEvent, endSSE } from './per-request.js';
import type { FastifyReply } from 'fastify';

function mockReply(): FastifyReply {
  const chunks: string[] = [];
  return {
    raw: {
      writeHead: vi.fn(),
      write: vi.fn((chunk: string) => {
        chunks.push(chunk);
        return true;
      }),
      end: vi.fn(),
      _chunks: chunks,
    },
    hijack: vi.fn(),
  } as unknown as FastifyReply;
}

describe('per-request SSE helpers', () => {
  describe('startSSE', () => {
    it('sets correct SSE headers', () => {
      const reply = mockReply();
      startSSE(reply);
      expect(reply.raw.writeHead).toHaveBeenCalledWith(200, {
        'Content-Type': 'text/event-stream',
        'Cache-Control': 'no-cache',
        Connection: 'keep-alive',
        'X-Accel-Buffering': 'no',
      });
    });

    it('hijacks the reply', () => {
      const reply = mockReply();
      startSSE(reply);
      expect(reply.hijack).toHaveBeenCalled();
    });
  });

  describe('sendSSEData', () => {
    it('formats object data as JSON', () => {
      const reply = mockReply();
      sendSSEData(reply, { hello: true });
      expect(reply.raw.write).toHaveBeenCalledWith('data: {"hello":true}\n\n');
    });

    it('formats string data without double-encoding', () => {
      const reply = mockReply();
      sendSSEData(reply, 'plain text');
      expect(reply.raw.write).toHaveBeenCalledWith('data: plain text\n\n');
    });

    it('formats number data as JSON', () => {
      const reply = mockReply();
      sendSSEData(reply, 42);
      expect(reply.raw.write).toHaveBeenCalledWith('data: 42\n\n');
    });
  });

  describe('sendSSEEvent', () => {
    it('formats named event with object data', () => {
      const reply = mockReply();
      sendSSEEvent(reply, 'token', { text: 'hi' });
      expect(reply.raw.write).toHaveBeenCalledWith('event: token\ndata: {"text":"hi"}\n\n');
    });

    it('formats named event with string data', () => {
      const reply = mockReply();
      sendSSEEvent(reply, 'status', 'done');
      expect(reply.raw.write).toHaveBeenCalledWith('event: status\ndata: done\n\n');
    });
  });

  describe('endSSE', () => {
    it('sends the [DONE] sentinel', () => {
      const reply = mockReply();
      endSSE(reply);
      expect(reply.raw.write).toHaveBeenCalledWith('data: [DONE]\n\n');
    });

    it('ends the stream', () => {
      const reply = mockReply();
      endSSE(reply);
      expect(reply.raw.end).toHaveBeenCalled();
    });
  });
});