File size: 2,052 Bytes
a42bc3b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | 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();
});
});
|