Testing AI code with pytest

On Thursday morning the central kitchen called Meera. Three customers had arrived at the Aundh store to collect 1.5 kg birthday cakes and found 1 kg cakes waiting for them. By lunch the count was eleven. By evening, going back through two days of WhatsApp orders, it was 37: every order for anything other than a whole kilogram or “half kg” had been sent to the kitchen as exactly 1 kg. Remakes, refunds and apology vouchers came to ₹18,400, in the week before Diwali.

The cause was Arjun’s change from Monday night. Customers kept writing “half kg”, which the old parser rejected, so he had changed the prompt and added a small conversion function. He tested it on three messages, all of which worked, and pushed it to production.

Meera did not raise her voice. “I don’t need to know whose fault it is. I need to know how we stop this happening again.”

“You can’t unit test an LLM”

Arjun’s first reaction, once the panic passed, was to write a test. He had written tests for the order database for years. He wrote this:

def test_parse_order():
    order = extract_order("2 chocolate truffle cakes, 1.5kg each, Aundh store", client)
    assert order.model_dump_json() == (
        '{"item":"chocolate truffle cake","count":2,"size":"1.5kg","store":"Aundh"}'
    )

It passed. He ran it again. It failed: the model had returned "Chocolate Truffle Cake" with capitals. He loosened it to compare lowercase, and the third run failed because the store came back as "Aundh store". Each run took four seconds and cost money. When he pushed the branch, the test failed in CI (continuous integration, the server that runs your tests on every change) because CI had no API key.

He came to Priya frustrated. “This is pointless. The model says something slightly different every time. You can’t unit test an LLM.”

“Correct,” Priya said. “So stop trying to unit test the LLM. Show me exactly what broke on Thursday.”

What actually broke

Arjun opened the diff. The prompt change asked the model to return the size “exactly as the customer wrote it”, so the model now returned strings like "half kg", "1.5kg" and "750 g". His new function converted those to kilograms:

def to_kg(size: str) -> float:
    if "half" in size:
        return 0.5
    if size.strip().endswith(" kg"):
        return float(size.split()[0])
    return 1.0   # sensible default for anything else

Priya read it once. “What does this return for 1.5kg, with no space?”

Arjun traced it in his head. No “half”. Does "1.5kg" end with " kg", with a space? No. So: the sensible default. 1.0.

“And 750 g? And 2 KG in capitals?”

“1.0. And 1.0.” He went quiet. “The model was doing its job perfectly. It returned exactly what the customer wrote. My four lines of plain Python turned every unusual size into one kilogram.”

“Which test would have caught that?”

“A boring one. Call to_kg with ten sizes and check the answers. No model needed.”

Priya drew a sandwich on the whiteboard. “Your app is a sandwich. The filling is the model call: unpredictable, a bit different every time. The bread is everything around it: building the prompt, parsing the output, converting units, checking stock, writing the kitchen ticket. The bread is ordinary code. It behaves the same way every time, so you can test it the ordinary way. Most of your bugs live in the bread.”

She added an analogy for Meera’s benefit. “It’s like baking. You can’t guarantee two cakes come out identical; ovens vary. But you can absolutely test that the scale reads 500 grams when there are 500 grams on it, and that the timer rings at 40 minutes. Thursday was a broken scale, not a bad oven.”

Where the analogy breaks: in a kitchen, the scale and the oven are separate machines. In an LLM app, the boundary between bread and filling is something you design. The more decisions you move out of the prompt and into code, the more of your app you can test precisely.

Following one order through the sandwich

Priya made Arjun trace one real Thursday order, step by step, marking each step as predictable or not.

The message: “2 chocolate truffle cakes 1.5kg each, Saturday 5pm pickup at Aundh”

Step 1, build the request (bread). The app combines the 60-token system prompt with the customer message, about 95 tokens in total. Same input, same request, every time. Predictable.

Step 2, call the model (filling). The model returns structured output, which the SDK checks against the Order shape:

{"item": "chocolate truffle cake", "count": 2, "size": "1.5kg", "store": "Aundh"}

About 1.5 seconds and well under a rupee. Run it again and the item might come back capitalised. Not predictable in its exact text.

Step 3, convert the size (bread). to_kg("1.5kg") returned 1.0. This is the bug. Predictable, and predictably wrong.

Step 4, write the kitchen ticket (bread). "2 x chocolate truffle cake @ 1 kg for Aundh". Predictable, faithfully passing on step 3’s mistake.

Four steps; one unpredictable. The bug was in a predictable one. Arjun had been trying to test the only step that cannot be tested precisely, and had skipped the three that can.

Opening the hood: four layers of tests

pytest is the standard Python testing tool: you write functions whose names start with test_, put assert statements in them, and run pytest to see which pass. For an LLM app, Priya sketched four layers, from fastest to slowest:

Layer	What it checks	Calls a real model?	Speed and cost	When it runs
Unit tests	Your plain functions: to_kg, prompt building, validation	No	Milliseconds, free	Every save, every commit
Mocked tests	Your code’s handling of model output, using a fake model	No	Milliseconds, free	Every commit
Snapshot tests	That prompts did not change by accident	No	Milliseconds, free	Every commit
Eval tests	That the real model still does the job well enough	Yes	Minutes, costs money	When prompts or models change

Two terms need defining. Mocking means replacing a real dependency, here the model client, with a fake object that returns whatever you tell it to, so you can test your code without the network, the cost or the randomness. A snapshot test saves a copy of something, such as a prompt, and fails if it ever changes without you deliberately updating the saved copy.

The bottom layer is the eval work from earlier in this module, packaged so it runs automatically: a regression eval, a fixed set of examples you re-run after every relevant change to make sure nothing that used to work has stopped working.

Building the tests, layer by layer

Here is the app code, with the conversion function rewritten properly. The important design choice is in extract_order: the client is passed in as an argument, rather than created inside the function. This is called dependency injection, and it is what makes the fake model possible.

# crumb_orders.py
import re
import anthropic
from pydantic import BaseModel

SYSTEM_PROMPT = """You turn WhatsApp messages into Crumb bakery orders.
Return the item, how many, the size exactly as the customer wrote it, and the store."""

class Order(BaseModel):
    item: str
    count: int
    size: str     # as written: "1.5kg", "half kg", "750 g"
    store: str

WORDS = {"half": 0.5, "quarter": 0.25, "one": 1.0, "two": 2.0}

def to_kg(size: str) -> float:
    """Convert '1.5kg', '750 g' or 'half kg' to kilograms. Unknown sizes raise."""
    text = size.lower().strip()
    match = re.match(r"([\d.]+)\s*(kg|g)", text)
    if match:
        number, unit = float(match.group(1)), match.group(2)
        return number if unit == "kg" else number / 1000
    for word, value in WORDS.items():
        if text.startswith(word):
            return value
    raise ValueError(f"Unknown size: {size!r}")   # never guess a cake size

def extract_order(whatsapp_text: str, client: anthropic.Anthropic) -> Order:
    response = client.messages.parse(
        model="claude-opus-5-5",
        max_tokens=1000,
        system=SYSTEM_PROMPT,
        messages=[{"role": "user", "content": whatsapp_text}],
        output_format=Order,
    )
    return response.parsed_output

def kitchen_ticket(order: Order) -> str:
    return f"{order.count} x {order.item} @ {to_kg(order.size):g} kg for {order.store}"

Notice the last line of to_kg. The old version had a “sensible default” of 1 kg. The new one refuses to guess: an unknown size raises an error, and the order goes to a human. A silent default was the real design mistake behind Thursday.

Layer 1: unit tests. Install pytest as a development dependency with uv add --dev pytest. Then:

# tests/test_units.py
import pytest
from crumb_orders import to_kg

@pytest.mark.parametrize("size, expected", [
    ("1.5 kg", 1.5),
    ("1.5kg", 1.5),      # no space: Thursday's bug
    ("half kg", 0.5),
    ("750 g", 0.75),
    ("2 KG", 2.0),       # capitals
])
def test_to_kg(size, expected):
    assert to_kg(size) == pytest.approx(expected)

def test_to_kg_refuses_to_guess():
    with pytest.raises(ValueError):
        to_kg("a big one")

@pytest.mark.parametrize runs the same test once per row, so five sizes become five tests. pytest.approx compares decimals safely (0.1 + 0.2 is not exactly 0.3 in computer arithmetic).

Run against Thursday’s buggy to_kg, with uv run pytest tests/test_units.py:

FAILED tests/test_units.py::test_to_kg[1.5kg-1.5] - assert 1.0 == 1.5 ± 1.5e-06
FAILED tests/test_units.py::test_to_kg[750 g-0.75] - assert 1.0 == 0.75 ± 7.5e-07
FAILED tests/test_units.py::test_to_kg[2 KG-2.0] - assert 1.0 == 2.0 ± 2.0e-06
FAILED tests/test_units.py::test_to_kg_refuses_to_guess - Failed: DID NOT RAISE
4 failed, 2 passed in 0.03s

Thirty milliseconds. No API key. No cost. Thursday’s bug, caught four ways. Against the fixed version: 6 passed.

Layer 2: mocked tests. Now test the code that talks to the model, using a fake client:

# tests/test_mocked.py
from types import SimpleNamespace
from unittest.mock import MagicMock
from crumb_orders import Order, SYSTEM_PROMPT, extract_order, kitchen_ticket

def fake_client(order: Order) -> MagicMock:
    """A stand-in for anthropic.Anthropic() that always returns `order`."""
    client = MagicMock()
    client.messages.parse.return_value = SimpleNamespace(parsed_output=order)
    return client

TRUFFLE = Order(item="chocolate truffle cake", count=2, size="1.5kg", store="Aundh")

def test_extract_order_sends_prompt_and_returns_order():
    client = fake_client(TRUFFLE)
    result = extract_order("2 truffle cakes 1.5kg each, Aundh", client)
    assert result == TRUFFLE
    sent = client.messages.parse.call_args.kwargs     # what our code sent
    assert sent["system"] == SYSTEM_PROMPT
    assert sent["output_format"] is Order

def test_ticket_from_model_output():
    assert kitchen_ticket(TRUFFLE) == "2 x chocolate truffle cake @ 1.5 kg for Aundh"

MagicMock creates an object that accepts any method call and records it. Setting return_value decides what the fake “model” answers. call_args.kwargs lets you check what your code sent. Expected result: 2 passed in 0.02s.

Layer 3: a snapshot test for the prompt. The prompt is code that changes behaviour. A one-word edit can shift results, and on Monday nobody noticed the prompt had changed at all. A snapshot test forces every prompt change to be deliberate:

# tests/test_snapshot.py
from pathlib import Path
import pytest
from crumb_orders import SYSTEM_PROMPT

SNAPSHOT = Path(__file__).parent / "snapshots" / "system_prompt.txt"

def test_system_prompt_unchanged():
    if not SNAPSHOT.exists():                    # first run: save the snapshot
        SNAPSHOT.parent.mkdir(exist_ok=True)
        SNAPSHOT.write_text(SYSTEM_PROMPT)
        pytest.skip("Snapshot created; commit it.")
    assert SYSTEM_PROMPT == SNAPSHOT.read_text(), (
        "Prompt changed. Run the eval suite (pytest -m eval), "
        "then delete the snapshot file to accept the new prompt."
    )

The test does not judge whether the new prompt is better. It makes sure a prompt change cannot slip into production without someone running the evals and updating the snapshot, which shows up clearly in code review.

Layer 4: an eval test with a real model. Arjun collected 40 real WhatsApp orders from the last month, including all the awkward sizes, with the correct count and kilograms for each, in tests/order_cases.jsonl:

# tests/test_evals.py
import json
import os
from pathlib import Path
import anthropic
import pytest
from dotenv import load_dotenv
from crumb_orders import extract_order, to_kg

load_dotenv()
CASES_FILE = Path(__file__).parent / "order_cases.jsonl"
CASES = [json.loads(line) for line in CASES_FILE.read_text().splitlines()]

@pytest.mark.eval
@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="needs an API key")
def test_order_extraction_pass_rate():
    client = anthropic.Anthropic()
    failures = []
    for case in CASES:
        order = extract_order(case["message"], client)
        try:
            ok = order.count == case["count"] and to_kg(order.size) == pytest.approx(case["kg"])
        except ValueError:
            ok = False                           # unknown size counts as a failure
        if not ok:
            failures.append((case["message"], order.model_dump()))
    pass_rate = 1 - len(failures) / len(CASES)
    print(f"\nPass rate: {pass_rate:.0%} ({len(CASES)} cases)")
    assert pass_rate >= 0.95, f"Below threshold. First failures: {failures[:3]}"

This test does not demand that every case passes. It demands a pass rate (the share of cases that succeed) of at least 95%, the binary-eval habit from earlier in this module. Each run makes 40 model calls, takes about a minute and costs roughly ₹12.

Register the marker so pytest knows about it, in pyproject.toml:

[tool.pytest.ini_options]
markers = ["eval: slow tests that call a real model"]

Now you have two commands. uv run pytest -m "not eval" runs everything fast and free, in under a second. uv run pytest -m eval runs the real-model evals.

Running it all in CI. A minimal GitHub Actions workflow that runs the fast tests on every pull request and the evals with a secret API key:

# .github/workflows/tests.yml
name: tests
on: [pull_request]
jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: astral-sh/setup-uv@v5
      - run: uv run pytest -m "not eval"
      - run: uv run pytest -m eval
        env:
          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}

The API key lives in the repository’s encrypted secrets, never in the code. Many teams run the eval step only when files under prompts/ or the model settings change, to save money.

The flaky eval on Tuesday

A week later the eval test failed on a pull request that only changed a README. Arjun re-ran it, and it passed. He re-ran it again: pass. Then fail.

That is a flaky test: one that sometimes passes and sometimes fails with no change to the code. It is poison, because after a few false alarms people stop believing failures, including the real ones.

Arjun looked at the numbers. The pass rate wobbled between 93% and 98% from run to run: on 40 cases, that is the difference between 2 and 3 failures. His threshold of 95% sat right in the middle of the wobble. Two cases were genuinely ambiguous (“1 and half kg” was read sometimes as 1.5 and sometimes as 1), and they flipped back and forth.

Priya’s checklist for taming flaky AI tests:

• Set the threshold below normal wobble. Run the eval ten times on a known-good version and look at the range. Crumb’s was 93% to 98%, so the threshold became 90%. A real regression, like Thursday’s, would have scored around 60%.
• Add more cases where you can afford it. Wobble shrinks as the test set grows. Going from 40 to 100 cases roughly halves the noise.
• Assert properties, not exact text. Check count == 2 and kg == 1.5, never the exact JSON string. Normalise case and whitespace before comparing.
• Fix ambiguous cases instead of tolerating them. “1 and half kg” was a specification gap; one example in the prompt made it consistent.
• Retry only infrastructure errors. Catch timeouts and rate limits and retry those, as you learned in Module 19. Never retry a quality failure.
• Record and replay where useful. Save real model responses once, then replay them in fast tests. Good for testing parsing against realistic output; useless for detecting model changes.

Where this approach breaks

Mocks can drift from reality. Six months later the SDK changes the shape of a response. Your fake client still returns the old shape, so every mocked test passes while production breaks. Keep one small “smoke test” that makes a single real call and checks the shape, run in the eval step.

Snapshot tests become rubber stamps. If the team deletes the snapshot every time it fails without running evals, it protects nothing. Make “eval results attached” part of reviewing any prompt change.

Evals in CI cost real money. At ₹12 per run and 30 pull requests a week, that is about ₹360 a week for one eval file. Fine for Crumb; with a 500-case suite and a judge model it gets expensive fast. Run the cheap layers always and the expensive layer selectively.

Tests are not monitoring. A test set of 40 cases cannot contain next month’s surprises. The tracing and error analysis earlier in this module catch what tests cannot imagine.

Try it yourself

Take any small LLM feature you built earlier in the course (a classifier or extractor from Module 04 works well).

1. Move one decision out of the prompt into a plain function, and write five parametrised unit tests for it, including at least one awkward input.
2. Make your model call take the client as an argument, and write one mocked test that checks what your code sends.
3. Add a snapshot test for the system prompt.
4. Write an eval test over 20 real examples with a pass-rate threshold. Run it five times and note the range of pass rates before choosing the threshold.

Expected result: uv run pytest -m "not eval" finishes in under a second with all tests passing, and your eval test’s pass rate varies by a few points between runs. Set the threshold below that range.

Common mistakes

• Asserting exact model output. It will fail randomly. Assert properties of parsed, normalised output.
• Creating the client inside your functions. It makes faking the model awkward. Pass the client in.
• Silent defaults for unknown input. A “sensible default” is a guess the rest of your system trusts. Raise an error and route to a human.
• Running real-model tests on every save. They are slow and cost money. Mark them and run them deliberately.
• Re-running flaky quality tests until they pass. It hides real regressions. Fix the threshold, the sample size or the ambiguous case.
• Thinking mocks test the model. They test your code. Evals test the model.

Recap

• An LLM app is a sandwich: predictable code around an unpredictable model call. Most bugs live in the predictable part.
• Unit-test your plain functions with pytest; they run in milliseconds with no API key.
• Pass the model client in, so tests can swap in a fake one with MagicMock.
• Snapshot your prompts so prompt changes are always deliberate and reviewed.
• Package your eval set as a marked pytest test with a pass-rate threshold, run in CI when prompts or models change.
• Tame flaky evals with thresholds below normal wobble, more cases and property-based assertions; retry only network errors.

On Friday, Arjun walked Meera through the pull request. “Thursday’s bug was in four lines of ordinary code, not in the AI. The AI read ‘1.5kg’ perfectly; my code turned it into 1 kg. Now there are 30 small tests that check that kind of code in a fraction of a second, every time anyone changes anything. If someone edits the AI’s instructions, the change is flagged and 40 real past orders are re-run to make sure we still get at least 90% right. And if the code ever sees a size it doesn’t understand, it asks a human instead of guessing.”

Meera nodded slowly. “So it can’t break by accident like that again.” Then she frowned. “What about on purpose? The support bot can issue refunds now. Could someone trick it into giving away cakes?”

Priya put down her coffee. “That’s next module. And yes, they could.”

Go deeper

• pytest documentation
• unittest.mock in the Python docs
• uv documentation
• AI Evals: Everything You Need to Know (Hamel Husain and Shreya Shankar)