Tests for prompts, tools and agent loops

The model's output is not deterministic, which leads people to conclude that AI systems cannot be unit tested. Most of the system is deterministic: the loop that runs tools, the budget that stops it, the parser that reads its output, the code that assembles the prompt. This section tests all of that by replacing the model with a fake that returns a script you wrote, so you can drive the loop through failure, repetition and exhaustion without spending anything or waiting on a network.

A fake that follows a script#

The agent loop has real logic: run tools, stop on budget, handle errors, detect repeats. To test it, replace the model with a fake that returns a scripted sequence of responses.

class ScriptedModel:
    def __init__(self, responses):
        self.responses = list(responses)
        self.calls = []

    def generate(self, messages, tools):
        self.calls.append(messages)
        return self.responses.pop(0)

def test_agent_runs_tool_then_answers():
    model = ScriptedModel([
        ToolCall("lookup_order", {"id": "A-102"}),
        FinalAnswer("Your order shipped yesterday."),
    ])
    result = run_agent("Where is my order?", model=model, tools=FAKE_TOOLS)
    assert result.text == "Your order shipped yesterday."
    assert "A-102" in str(model.calls[1])     # the tool result was sent back

Note what the fake records as well as what it returns. Keeping the messages it was called with is what lets you assert on the thing nobody checks: that the tool result actually made it back into the conversation. A fake that only returns is half a fake.

Test the edges of the loop#

• The same tool call repeated three times triggers the repeat guard.
• A tool that raises produces an error result, not a crash.
• Hitting the turn budget returns a partial result with a reason.
• An unknown tool name is rejected before anything runs.

Mocking a model call properly#

Two mistakes make a mocked model call worthless, and both produce a green test. The first is patching the name where the class was defined rather than the name the code under test actually reads. Python binds names at import time, so a module that did from app.provider import Client holds its own reference, and patching app.provider.Client afterwards changes nothing it will ever look at.

# app/summarise.py
from app.provider import Client
client = Client()

def summarise(text):
    return client.chat([{"role": "user", "content": "Summarise: " + text}]).strip()


# the test
def test_patching_where_it_was_defined():
    with patch("app.provider.Client") as fake:          # wrong target
        fake.return_value.chat.return_value = "  short  "
        assert summarise("long text") == "short"

def test_patching_where_it_is_looked_up():
    with patch("app.summarise.client") as fake:         # right target
        fake.chat.return_value = "  short  "
        assert summarise("long text") == "short"

Prove to yourself that the first one does nothing by making the real client refuse to work. Give it a chat method that raises, and the wrong patch fails loudly instead of quietly reaching the network.

$ pytest test_patch_target.py -q
app/summarise.py:7: in summarise
    reply = client.chat([{"role": "user", "content": "Summarise: " + text}])
E       RuntimeError: real network call: this must never run in a test
1 failed, 1 passed

The second mistake is a plain MagicMock, which agrees with everything. It will accept a method that does not exist, arguments in the wrong order and keyword arguments you invented, and every assertion about it will pass. create_autospec builds the mock from the real class, so a typo or a changed signature fails the test instead of sailing through it.

from unittest.mock import MagicMock, create_autospec

def test_a_plain_mock_agrees_with_anything():
    fake = MagicMock()
    fake.chatt(messages=[], temperature="hot", nonsense=1)   # no such method
    assert fake.chatt.called            # passes, and proves nothing

def test_autospec_rejects_the_typo():
    fake = create_autospec(Client, instance=True)
    fake.chatt([])                      # AttributeError, as it should be

For the provider SDK specifically, prefer a small hand-written fake class over patching anything. It is the same amount of code, it documents the shape your system depends on in one readable place, and when the SDK changes your fake stops matching, which is information. A mock silently keeps agreeing with a signature that no longer exists.

Recorded responses#

For an integration test that should exercise the real API format, record one real response to a file and replay it in tests. Libraries such as vcrpy do this for HTTP. The test is fast, free and repeatable, and you rerecord deliberately when the API changes.

Two rules keep recordings from becoming a liability. Scrub the API key out of the recorded request headers before it is committed, because a recording is a verbatim copy of a real exchange. And put a date in the file and rerecord on a schedule, otherwise the suite is testing an API shape from eighteen months ago and passing confidently.

Let the computer invent inputs#

Property-based testing generates many random inputs and checks that a property always holds. It is excellent for parsers of model output, which meet strange input in production.

from hypothesis import given, strategies as st

@given(st.text())
def test_parser_never_crashes(raw):
    result = parse_model_output(raw)          # must return a result or a clean error
    assert result.ok or result.error

What makes it worth the import is shrinking. When it finds a failure it keeps simplifying the input until it has the smallest one that still breaks, so you are handed a minimal reproduction rather than a page of noise. Pointed at the parser everybody writes first, it reports the answer in one line.

def classify_naive(raw):
    """The version everyone writes first."""
    return json.loads(raw[raw.find("{") : raw.rfind("}") + 1])["label"]

@given(st.text())
def test_naive_parser_never_raises(raw):
    assert classify_naive(raw) in LABELS

Snapshot the prompt#

Prompt assembly code, which combines templates, retrieved text and history, breaks quietly. A snapshot test stores the assembled prompt for a fixed input and fails when it changes, so every change to the prompt is reviewed on purpose.

Snapshot the prompt you send, never the answer you get back. The prompt is your code and must be stable; the answer is the model's and never will be. And keep the snapshot small enough to read in a diff, because a snapshot nobody reads gets accepted with the update flag on reflex, which is the same as not having it.

What a flaky test is telling you#

A flaky test is not noise to be retried. It is a report that something in your system depends on an order, a clock or a machine that nobody declared, and most of the time the bug is in the code rather than the test. Here is one that is not obviously random at all.

def top_tags(docs):
    """Returns the distinct tags. The set is an implementation detail."""
    return list({t for d in docs for t in d["tags"]})

def test_top_tags():
    docs = [{"tags": ["rag", "eval", "prompt"]}]
    assert top_tags(docs) == ["rag", "eval", "prompt"]

$ for i in 1 2 3 4 5; do pytest test_flaky_set.py -q | tail -1; done
1 failed in 0.13s
1 passed in 0.13s
1 failed in 0.15s
1 failed in 0.14s
1 failed in 0.13s

E  AssertionError: assert ['eval', 'prompt', 'rag'] == ['rag', 'eval', 'prompt']

The instructive part is what it found. The test is wrong to assert an order, but the function is also wrong to return a list: it promises an order it does not have, and somewhere downstream a caller is showing the first three tags to a user and getting a different three every deploy. Fix both. Return a sorted list or a set, and assert on the same thing.

Flakes when                           What it is actually telling you
run in a different order              tests share state, or a fixture is mutable
run in parallel but not alone         shared file, port, database row or temp path
only in CI                            a clock, a locale, a timezone, or less memory
about once in twenty runs             a real race, or a set or dict order dependency
after the first of the month          a date fixture built from datetime.now
only on a colleague's machine         an undeclared dependency on the environment

What never to do is add a retry decorator to make it green. A retry hides the failure rate rather than removing it, and the underlying race is still there in production where nothing reruns. If a test genuinely cannot be made deterministic, delete it and write an eval instead, because that is what it was.

What to skip#

• Testing that the model is any good. Different job, different tool, different schedule.
• Mocking the HTTP layer by hand. Record it, or fake the client above it.
• Asserting on exact prompt text inside a unit test. Snapshot it once, in one place.
• Property-based tests on business logic. Keep them for parsers and validators, where the input space is genuinely unbounded.