Asserting on a Stream of Chunks Instead of a Final String

#ai #machinelearning #tutorial #llm

Joining every delta and comparing the result to an expected sentence is the assertion most streaming tests make, and it is testing the model rather than your code. Everything you actually control lives in the structure of the chunks, which the join throws away.

Why the joined string is the wrong subject

Two problems, and they pull in opposite directions. The joined string is too strict about the thing you do not control — a model can reword an answer between runs, between versions, and at a non-zero temperature within a single version, so an exact comparison is a test that fails for reasons unrelated to your change. And it is too loose about the thing you do control: a stream that delivered every token in one chunk, in the wrong order, with a duplicated first delta, with no terminal frame, or after buffering for nine seconds, joins to exactly the same string as a healthy one.

Softening the exact comparison to toContain or a regex does not fix either half. It fixes the flakiness by asserting almost nothing, and it still says nothing about ordering, framing or termination. The way out is to stop asserting on the text and start asserting on the structure the text arrived in.

A collector that records structure

Write one helper that consumes a stream and returns everything an assertion might want. Doing this once means the assertions below are each a single line, which is what makes writing several of them bearable.

type Collected = {
  events: any[];        // every parsed data payload, in arrival order
  text: string;         // the concatenation, for the few cases that need it
  deltaCount: number;   // frames that carried content
  finishReasons: (string | null)[];
  sawTerminal: boolean;
};

async function collect(res: Response): Promise<Collected> {
  const out: Collected = {
    events: [], text: "", deltaCount: 0, finishReasons: [], sawTerminal: false,
  };
  let buffer = "";
  const reader = res.body!.pipeThrough(new TextDecoderStream()).getReader();
  for (;;) {
    const { value, done } = await reader.read();
    if (done) break;
    buffer += value;
    let i: number;
    while ((i = buffer.indexOf("\n\n")) !== -1) {
      const frame = buffer.slice(0, i);
      buffer = buffer.slice(i + 2);
      for (const line of frame.split("\n")) {
        if (!line.startsWith("data:")) continue;      // skip ": " comments and event: lines
        const payload = line.slice(5).trimStart();
        if (payload === "[DONE]") { out.sawTerminal = true; continue; }
        const event = JSON.parse(payload);
        out.events.push(event);
        const choice = event.choices?.[0];
        if (typeof choice?.delta?.content === "string" && choice.delta.content !== "") {
          out.text += choice.delta.content;
          out.deltaCount++;
        }
        if (choice) out.finishReasons.push(choice.finish_reason);
      }
    }
  }
  return out;
}

The payload === "[DONE]" branch has to come before JSON.parse. OpenAI’s terminal line is the literal text [DONE] and is not valid JSON for that schema, so a collector that parses first throws on the last frame of every successful stream. Anthropic’s Messages API has no such sentinel at all, so the same collector needs a different terminal condition there — its named message_stop event.

Shape assertions

The strongest assertion you can make about a chunk is that it matches the schema the provider documents, because that holds for every prompt and catches a provider changing its wire format. Validate every event, not the first one.

import { z } from "zod";

const Chunk = z.object({
  id: z.string(),
  object: z.literal("chat.completion.chunk"),
  model: z.string(),
  choices: z.array(
    z.object({
      index: z.number().int(),
      delta: z.object({
        role: z.literal("assistant").optional(),
        content: z.string().nullable().optional(),
        tool_calls: z.array(z.any()).optional(),
      }),
      finish_reason: z
        .enum(["stop", "length", "tool_calls", "content_filter", "function_call"])
        .nullable(),
    }),
  ),
  usage: z.any().nullable().optional(),
});

test("every chunk matches the documented shape", async () => {
  const { events } = await collect(await callEndpoint("hello"));
  for (const [i, event] of events.entries()) {
    const parsed = Chunk.safeParse(event);
    expect(parsed.success, "chunk " + i + ": " + JSON.stringify(event)).toBe(true);
  }
});

Note that the enum lists every documented finish_reason value rather than only stop. A schema that admits only the happy value turns a legitimate length truncation into a schema failure, and you will spend an afternoon on it. Include the index in the failure message: with sixty chunks, “expected true, got false” on its own is nearly useless.

Ordering assertions

Ordering is where the real invariants are, because these are properties the provider guarantees and your proxy can break. For an OpenAI-shaped stream:

delta.role appears on the first chunk and nowhere else. Assert events.filter((e) => e.choices[0].delta.role).length is 1 and that the match is at index 0.
Exactly one non-null finish_reason, and it is on the last chunk that has a choices entry. Assert finishReasons.filter(Boolean) has length 1.
No content delta arrives after the chunk carrying finish_reason. This is the assertion that catches a proxy reordering or re-emitting frames.
With stream_options: {"include_usage": true}, usage arrives on one additional chunk whose choices array is empty; every earlier chunk carries usage: null. Assert that the usage-bearing chunk is last and that no chunk has both content and usage.

Tool calls add one more ordering rule that is easy to miss. When a model calls tools while streaming, the deltas carry a tool_calls array whose entries have an index, and only the first delta for a given index carries the call id and the function name; every later delta for that index carries a fragment of function.arguments and nothing else. So the invariant is: for each index, exactly one delta supplies a name, it is the first, and no delta supplies a name twice. An accumulator that overwrites rather than concatenates passes a text-only test and fails here, which is why this assertion belongs in the same file rather than in the tool-calling suite.

For an Anthropic-shaped stream the equivalent invariants are about event names: a content_block_start precedes any delta at the same index, every started block is closed by a content_block_stop, and message_stop is last. A useful single assertion is that the sequence of event types, with ping filtered out, matches a regular expression over the type names — that expresses the whole grammar in one line and fails with a readable diff.

Counting, carefully

Chunk count is tempting and mostly a trap. How many chunks a provider emits for a given answer depends on tokenisation, batching inside the inference server and network coalescing; it is not part of any published contract, and asserting deltaCount === 14 gives you a test that fails on a good day. What you can assert:

Against a fixture, an exact count is fine. When the upstream is your own ReadableStream, the count is fully determined and asserting it catches your proxy merging or splitting frames.
Against a real provider, assert a lower bound. expect(deltaCount).toBeGreaterThan(1) is the assertion that distinguishes streaming from not streaming, and it is stable.
Assert the relationship, not the number. text equals the concatenation of the deltas in order, and equals the non-streamed response for the same request at temperature zero. That second one is a genuine invariant of the two endpoints and far more valuable than any count.

The same reasoning applies to token totals: compare the usage figure the stream reports against your own count only as an approximate check, because the two are produced by different code paths and can legitimately differ — see token count mismatch for why. What you assert is that a usage figure arrived at all, exactly once, on the frame the provider documents it on.

One assertion that looks like counting but is not: incrementality over time. Record the arrival timestamp of each frame alongside its content and assert that the first frame arrived meaningfully before the last. Against a fixture with a deliberate gap between enqueues this is deterministic, and it is the only assertion in the set that fails when a proxy collects the whole response and flushes it once. It costs a single array and it is the difference between testing that your endpoint returns the right bytes and testing that it streams them.