1
import { describe, expect, it } from "vitest";
2
import { createMockAiClient } from "@/worker/lib/ai/mock";
3
import type { AiFrame } from "@/worker/lib/ai/types";
4
 
5
async function collectChunks(iter: AsyncIterable<AiFrame>): Promise<string> {
6
  let buffer = "";
7
  for await (const frame of iter) {
8
    if (frame.type === "chunk") buffer += frame.text;
9
  }
10
  return buffer;
11
}
12
 
13
async function collectFrames(iter: AsyncIterable<AiFrame>): Promise<AiFrame[]> {
14
  const frames: AiFrame[] = [];
15
  for await (const frame of iter) frames.push(frame);
16
  return frames;
17
}
18
 
19
describe("mock AI client", () => {
20
  it("emits a deterministic [mock-chat] prefix from the last user message", async () => {
21
    const client = createMockAiClient();
22
    const iter = await client.chat([
23
      { role: "system", content: "You are an assistant." },
24
      { role: "user", content: "Please proofread this sentence." },
25
    ]);
26
    const body = await collectChunks(iter);
27
    expect(body.startsWith("[mock-chat]")).toBe(true);
28
    expect(body).toContain("Please proofread this sentence.");
29
  });
30
 
31
  it("summarize returns a deterministic [mock-summary] prefix with usage", async () => {
32
    const client = createMockAiClient();
33
    const result = await client.summarize("Launch freeze begins Thursday. Rollout staged by region.");
34
    expect(result.summary.startsWith("[mock-summary]")).toBe(true);
35
    expect(result.summary).toContain("Launch freeze begins Thursday.");
36
    expect(result.usage).toEqual({ promptTokens: 10, completionTokens: 5, totalTokens: 15 });
37
  });
38
 
39
  it("chat stream emits a usage frame at the end", async () => {
40
    const client = createMockAiClient();
41
    const iter = await client.chat([{ role: "user", content: "hi" }]);
42
    const frames = await collectFrames(iter);
43
    const usage = frames.find((f): f is Extract<AiFrame, { type: "usage" }> => f.type === "usage");
44
    expect(usage?.usage).toEqual({ promptTokens: 10, completionTokens: 5, totalTokens: 15 });
45
    expect(frames[frames.length - 1]?.type).toBe("usage");
46
    expect(frames.filter((f) => f.type === "chunk").length).toBeGreaterThan(0);
47
  });
48
});