58 lines
2.2 KiB
TypeScript
58 lines
2.2 KiB
TypeScript
// 不自动运行（需要 npm i -D vitest vue 后执行 npx vitest run）
|
||
import { describe, expect, it } from "vitest";
|
||
import { effectScope } from "vue";
|
||
import { useChatStream } from "./useChatStream.ts";
|
||
|
||
function sseResponse(chunks: string[], failAfter?: number): Response {
|
||
let i = 0;
|
||
const body = new ReadableStream<Uint8Array>({
|
||
pull(ctrl) {
|
||
if (failAfter !== undefined && i === failAfter) return ctrl.error(new Error("network down"));
|
||
if (i >= chunks.length) return ctrl.close();
|
||
ctrl.enqueue(new TextEncoder().encode(chunks[i++]));
|
||
},
|
||
});
|
||
return new Response(body, { status: 200, headers: { "Content-Type": "text/event-stream" } });
|
||
}
|
||
|
||
const FINAL = 'id: 3\nevent: final\ndata: {"answer":"好","citations":[{"id":1,"title":"政策"}]}\n\n';
|
||
|
||
describe("useChatStream", () => {
|
||
it("拼接 token 并解析 final", async () => {
|
||
const scope = effectScope();
|
||
const chat = scope.run(() =>
|
||
useChatStream({ url: "/chat", fetchImpl: async () => sseResponse(["id: 1\nevent: token\ndata: 你", "好\n\n", FINAL]) }),
|
||
)!;
|
||
await chat.send("hi");
|
||
expect(chat.status.value).toBe("done");
|
||
expect(chat.text.value).toBe("你好");
|
||
expect(chat.final.value?.citations[0]?.title).toBe("政策");
|
||
scope.stop();
|
||
});
|
||
|
||
it("断线后带 Last-Event-ID 重试", async () => {
|
||
const seen: (string | null)[] = [];
|
||
let n = 0;
|
||
const fetchImpl = async (_: RequestInfo | URL, init?: RequestInit) => {
|
||
seen.push(new Headers(init?.headers).get("Last-Event-ID"));
|
||
return n++ === 0 ? sseResponse(["id: 1\nevent: token\ndata: 你\n\n", "x"], 1) : sseResponse([FINAL]);
|
||
};
|
||
const chat = useChatStream({ url: "/chat", fetchImpl, maxRetries: 1 });
|
||
await chat.send("hi");
|
||
expect(seen).toEqual([null, "1"]);
|
||
expect(chat.status.value).toBe("done");
|
||
});
|
||
|
||
it("stop 之后状态为 aborted", async () => {
|
||
const fetchImpl = (_: RequestInfo | URL, init?: RequestInit) =>
|
||
new Promise<Response>((_, reject) =>
|
||
init?.signal?.addEventListener("abort", () => reject(new DOMException("aborted", "AbortError"))),
|
||
);
|
||
const chat = useChatStream({ url: "/chat", fetchImpl });
|
||
const p = chat.send("hi");
|
||
chat.stop();
|
||
await p;
|
||
expect(chat.status.value).toBe("aborted");
|
||
});
|
||
});
|