Files
Genarrative/src/services/sseStream.test.ts

99 lines
2.5 KiB
TypeScript

import { expect, test } from 'vitest';
import { readSseJsonStream, readSseStream } from './sseStream';
function createChunkedStreamResponse(chunks: Uint8Array[]) {
const stream = new ReadableStream<Uint8Array>({
start(controller) {
for (const chunk of chunks) {
controller.enqueue(chunk);
}
controller.close();
},
});
return new Response(stream, {
headers: {
'Content-Type': 'text/event-stream; charset=utf-8',
},
});
}
test('readSseJsonStream flushes decoder tail and handles CRLF boundaries', async () => {
const encoder = new TextEncoder();
const prefix = encoder.encode('event: reply_delta\r\ndata: {"text":"');
const replyBytes = encoder.encode('溪上春风');
const suffix = encoder.encode('"}\r\n\r\n');
const splitIndex = replyBytes.length - 1;
const events: Array<{ eventName: string; parsed: Record<string, unknown> }> =
[];
await readSseJsonStream(
createChunkedStreamResponse([
new Uint8Array([...prefix, ...replyBytes.slice(0, splitIndex)]),
new Uint8Array([...replyBytes.slice(splitIndex), ...suffix]),
]),
({ eventName, parsed }) => {
events.push({ eventName, parsed });
},
);
expect(events).toEqual([
{
eventName: 'reply_delta',
parsed: { text: '溪上春风' },
},
]);
});
test('readSseJsonStream skips malformed json and keeps valid LF events', async () => {
const encoder = new TextEncoder();
const events: Array<{ eventName: string; parsed: Record<string, unknown> }> =
[];
await readSseJsonStream(
createChunkedStreamResponse([
encoder.encode(
'event: malformed\ndata: not-json\n\n' +
'event: ready\ndata: {"value":7}\n\n',
),
]),
({ eventName, parsed }) => {
events.push({ eventName, parsed });
},
);
expect(events).toEqual([
{
eventName: 'ready',
parsed: { value: 7 },
},
]);
});
test('readSseStream can stop early and cancel the reader', async () => {
const encoder = new TextEncoder();
let cancelled = false;
const stream = new ReadableStream<Uint8Array>({
start(controller) {
controller.enqueue(
encoder.encode(
'event: first\ndata: one\n\n' + 'event: second\ndata: two\n\n',
),
);
},
cancel() {
cancelled = true;
},
});
const events: string[] = [];
await readSseStream(new Response(stream), ({ eventName }) => {
events.push(eventName);
return false;
});
expect(events).toEqual(['first']);
expect(cancelled).toBe(true);
});