Test AI code with a fake model

Test the whole generation pipeline offline with a fake LangChain model that replays scripted replies: happy paths, repairs, 429 retries and fallbacks, in milliseconds.

Lesson 38 of 4414 min+40 XPHands-on

Testing AI code against the real model is slow (30 seconds a deck), costs money, and gives different output every run. Worse, the paths that matter most, like a 429 halfway through or an overloaded model, happen when they want, not when you test. A fixes all of that.

generateDeck(topic, { model })

createModel()

real Gemini: slow, costs tokens, varies

createFakeModel({ replies })

instant, free, same every run

generateDeck only needs something with invoke, stream and withStructuredOutput. In production that's Gemini; in tests it's a fake that replays a script.

What the pipeline needs from a model

Only four things. A fake that has them is, as far as the pipeline can tell, a real model:

invoke(messages): returns { content, usage_metadata }.
stream(messages): yields { content } chunks, with usage_metadata on the last one.
withStructuredOutput(schema, options): returns something with its own invoke.
bindTools(): returns the model.

Scripted replies

The fake takes a list of replies and hands them out in order, one per invoke or stream. A reply can be text, or an Error, which is thrown instead: apiError(429) makes the next call look rate-limited. It also records every call in calls, so a test can check what the pipeline sent.

test/pipeline.test.js (one test)

test("retries a rate-limited call, then succeeds", async () => {
  const model = createFakeModel({ outline: OUTLINE, replies: [apiError(429), GOOD_DECK] });
  const events = [];
  const deck = await generateDeck("cobots", { model, wait: noWait, onEvent: (event) => events.push(event) });

  assert.equal(deck.slides.length, 3);
  assert.equal(events.filter((event) => event.type === "retry").length, 1);
});

Your project so far

28 files · 2 new or changed in this lesson

test/pipeline.test.js

// Run with: node --test
// No API key, no network, no cost. Every test finishes in milliseconds.

import test from "node:test";
import assert from "node:assert/strict";
import { generateDeck } from "../src/pipeline.js";
import { apiError, createFakeModel } from "./fake-model.js";

const OUTLINE = {
  title: "Cobots",
  slides: [
    { component: "TitleSlide", heading: "Cobots", points: [] },
    { component: "BulletsSlide", heading: "Why now", points: ["Labor", "Cost"] },
    { component: "ClosingSlide", heading: "Start small", points: [] },
  ],
};

const GOOD_DECK = `<deck>
  <TitleSlide title="Cobots in Small Factories"/>
  <BulletsSlide title="Why now"><Bullet>Skilled labor is scarce</Bullet><Bullet>Payback under a year</Bullet></BulletsSlide>
  <ClosingSlide title="Start small"/>
</deck>`;

// Slide 2 forgets its required title.
const BROKEN_DECK = GOOD_DECK.replace('<BulletsSlide title="Why now">', "<BulletsSlide>");
const FIXED_SLIDE_2 = '<fix slide="2"><BulletsSlide title="Why now"><Bullet>Skilled labor is scarce</Bullet></BulletsSlide></fix>';

const noWait = async () => {};

test("builds a deck from scripted replies", async () => {
  const model = createFakeModel({ outline: OUTLINE, replies: [GOOD_DECK] });
  const deck = await generateDeck("cobots", { model, wait: noWait });

  assert.equal(deck.slides.length, 3);
  assert.deepEqual(deck.problems, []);
  assert.equal(deck.repairs, 0);
});

test("sends only the broken slide back for repair", async () => {
  const model = createFakeModel({ outline: OUTLINE, replies: [BROKEN_DECK, FIXED_SLIDE_2] });
  const events = [];
  const deck = await generateDeck("cobots", { model, wait: noWait, onEvent: (event) => events.push(event) });

  assert.deepEqual(deck.problems, []);
  assert.equal(deck.repairs, 1);
  const repairRequest = JSON.stringify(model.calls.at(-1).input);
  assert.match(repairRequest, /slide=\\"2\\"/);
  assert.doesNotMatch(repairRequest, /Start small/); // slide 3 was fine, so it wasn't sent
  assert.equal(events.filter((event) => event.type === "slide").length, 3); // slides streamed as they arrived
});

test("retries a rate-limited call, then succeeds", async () => {
  const model = createFakeModel({ outline: OUTLINE, replies: [apiError(429), GOOD_DECK] });
  const events = [];
  const deck = await generateDeck("cobots", { model, wait: noWait, onEvent: (event) => events.push(event) });

  assert.equal(deck.slides.length, 3);
  assert.equal(events.filter((event) => event.type === "retry").length, 1);
});

test("falls back to the second model when the first is overloaded", async () => {
  const busy = createFakeModel({ name: "busy", outline: OUTLINE, replies: [apiError(503), apiError(503), apiError(503)] });
  const spare = createFakeModel({ name: "spare", outline: OUTLINE, replies: [GOOD_DECK] });
  const deck = await generateDeck("cobots", { models: [busy, spare], wait: noWait });

  assert.equal(deck.slides.length, 3);
  assert.equal(spare.calls.filter((call) => call.kind === "stream").length, 1);
});

test("does not retry a used-up daily quota", async () => {
  const model = createFakeModel({ outline: OUTLINE, replies: [apiError(429, "Quota exceeded: limit: 0"), GOOD_DECK] });
  await assert.rejects(generateDeck("cobots", { model, wait: noWait }), /limit: 0/);
});

Key takeaways

Because the pipeline takes a model as a parameter, tests can hand it a fake one.
Scripted replies make tests fast, free and identical every run, even for errors and retries.
Test the paths that are hard to trigger for real: repairs, rate limits, overloaded models.

Sign in to run the exercise

Reading is free. Writing code here needs an account so we have somewhere to keep your Gemini key and the +40 XP you are about to earn.

Continue with Google