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 { 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(); }); });