frontend/src/ui/lib/streamChat.test.ts

import { describe, expect, it, vi } from "vitest";
import { parseSseBuffer, streamChat } from "./streamChat";

describe("parseSseBuffer", () => {
  it("splits complete frames and keeps the remainder", () => {
    const { frames, rest } = parseSseBuffer(
      'data: {"delta":"a"}\n\ndata: {"delta":"b"}\n\ndata: {"del',
    );
    expect(frames).toEqual([{ delta: "a" }, { delta: "b" }]);
    expect(rest).toBe('data: {"del');
  });
});

function sseStream(chunks: string[]): ReadableStream<Uint8Array> {
  const enc = new TextEncoder();
  let i = 0;
  return new ReadableStream({
    pull(controller) {
      if (i < chunks.length) controller.enqueue(enc.encode(chunks[i++]));
      else controller.close();
    },
  });
}

describe("streamChat", () => {
  it("delivers deltas in order and the final route info", async () => {
    vi.stubGlobal("fetch", vi.fn(async () =>
      new Response(
        sseStream([
          'data: {"delta":"Hi"}\n\n',
          'data: {"delta":" there"}\n\n',
          'data: {"done":true,"compute":{"target":"device:colab","label":"Colab T4","fell_back":false}}\n\n',
        ]),
        { status: 200, headers: { "content-type": "text/event-stream" } },
      ),
    ));

    const deltas: string[] = [];
    let compute: any = null;
    await streamChat("http://x", { messages: [{ role: "user", content: "hi" }] }, {
      onDelta: (t) => deltas.push(t),
      onDone: (c) => (compute = c),
    });

    expect(deltas.join("")).toBe("Hi there");
    expect(compute).toMatchObject({ label: "Colab T4" });
    vi.unstubAllGlobals();
  });

  it("reports an error frame", async () => {
    vi.stubGlobal("fetch", vi.fn(async () =>
      new Response(sseStream(['data: {"error":"cloud unreachable"}\n\n', 'data: {"done":true,"compute":null}\n\n']), {
        status: 200,
      }),
    ));
    const errors: string[] = [];
    await streamChat("http://x", { messages: [] }, { onDelta: () => {}, onError: (e) => errors.push(e) });
    expect(errors).toEqual(["cloud unreachable"]);
    vi.unstubAllGlobals();
  });
});