Unit testing LLM agents: a practical guide

Most teams bolt tests onto their LLM agents after the fact, then wonder why regressions slip through. Unit testing LLM agents demands a different approach than testing deterministic functions: isolate the model from the orchestration logic, mock the nondeterministic parts, and assert on structured contracts rather than prose.

1. Separate the agent loop from the model call

An agent is a state machine that calls a model, parses the response, and possibly invokes tools. The bug surface is in the state transitions, not in the model weights. Extract the model call behind an interface so you can swap in a fake.

class LLMClient:
    def complete(self, messages, tools=None, **kwargs):
        raise NotImplementedError

class Agent:
    def __init__(self, llm: LLMClient):
        self.llm = llm
        self.history = []

    def step(self, user_msg: str):
        self.history.append({"role": "user", "content": user_msg})
        resp = self.llm.complete(self.history, tools=TOOL_SPECS)
        # parse resp, maybe call tools, append assistant message
        return resp

Now Agent contains zero network code. That is the unit under test.

2. Mock the provider at the boundary

Use a fake that returns OpenAI-compatible chat completion objects. This keeps your test faithful to the real schema, including tool call structures.

from types import SimpleNamespace

class FakeLLM(LLMClient):
    def __init__(self, canned_responses):
        self.responses = canned_responses
        self.call_count = 0
        self.last_request = None

    def complete(self, messages, tools=None, **kwargs):
        self.last_request = {"messages": messages, "tools": tools}
        resp = self.responses[self.call_count]
        self.call_count += 1
        return resp

fake = FakeLLM([
    SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(
        content=None,
        tool_calls=[{"id": "call_1", "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}}]
    ))])
])

Pitfall: ignoring streaming and async variants. If your agent uses stream=True, mock the iterator too, or you will get green tests that break in prod.

3. Assert on tool calls, not text

When the agent uses tool calling, the critical logic is argument extraction and dispatch. Assert the fake received the right tools and that the agent executed the local function.

agent = Agent(fake)
result = agent.step("What's the weather in SF?")
assert fake.last_request["tools"] == TOOL_SPECS
assert result.tool_calls[0]["function"]["name"] == "get_weather"
assert json.loads(result.tool_calls[0]["function"]["arguments"])["city"] == "SF"

If your agent loops until no tool calls, drive multiple canned responses and verify the termination condition.

Common mistake: asserting on content strings

Model output text drifts. A test like assert "sunny" in result.content will rot. Test the decision (which tool, which args) and leave fluency to eval suites.

4. Test parser and schema validation

Many agents use JSON mode or Pydantic to force structure. Unit test the parsing layer with both valid and malformed inputs.

from pydantic import BaseModel, ValidationError

class WeatherArgs(BaseModel):
    city: str
    units: str = "metric"

def parse_args(raw: str) -> WeatherArgs:
    return WeatherArgs.model_validate_json(raw)

# valid
assert parse_args('{"city": "NYC"}').city == "NYC"
# invalid
try:
    parse_args('{"city": 123}')
    assert False, "should have raised"
except ValidationError:
    pass

Tradeoff: strict schemas reduce model freedom. You gain testability, lose some natural language flexibility. Pick the strictness your product needs.

5. Simulate provider failures and fallback

Network errors, 429s, and malformed streaming chunks are where agents die in production. Inject failures in the fake.

class FlakyFake(LLMClient):
    def complete(self, messages, tools=None, **kwargs):
        raise TimeoutError("provider deadline exceeded")

If you route through a gateway such as n4n.ai, automatic fallback when a provider is rate-limited or degraded masks upstream outages, but your unit tests should still inject simulated 429s to verify your agent degrades gracefully or retries with backoff. The gateway’s OpenAI-compatible endpoint for 240+ models does not relieve you of handling TimeoutError in your own loop.

Test the retry policy:

def test_agent_retries_on_timeout():
    flaky = FlakyFake()
    agent = Agent(flaky)
    with pytest.raises(TimeoutError):
        agent.step("hi")
    assert flaky.call_count == MAX_RETRIES

6. Record real responses as fixtures

For regression tests, capture sanitized production transcripts. Store the model’s raw completion (no PII) as a JSON fixture and replay it through the fake.

{
  "messages": [{"role": "user", "content": "book a flight"}],
  "tool_calls": [{"function": {"name": "book_flight", "arguments": "{\"dest\": \"LAX\"}"}}]
}

Load it in the test:

def test_regression_book_flight():
    fixture = json.load(open("fixtures/book_flight.json"))
    fake = FakeLLM([make_response(fixture)])
    agent = Agent(fake)
    out = agent.step("book a flight")
    assert out.tool_calls[0]["function"]["name"] == "book_flight"

Tradeoff: fixtures rot as your prompt changes. Treat them like snapshot tests—update deliberately, not on every run.

7. Property-based tests for prompt construction

The part of the agent you can test deterministically is how it assembles the context window. Use property-based testing to check invariants.

from hypothesis import given, strategies as st

@given(st.lists(st.text(), min_size=1))
def test_system_prompt_always_present(messages):
    builder = PromptBuilder(system="You are a bot")
    full = builder.build(messages)
    assert full[0]["role"] == "system"
    assert len(full) == len(messages) + 1

This catches accidental truncation of few-shot examples or metadata.

8. Measure coverage on the right code

Do not aim for line coverage inside the LLM. Aim for coverage of your tool dispatch, error handling, and state transitions. A coverage report that excludes the llm.complete call site is honest.

pytest --cov=agent.core --cov-report=term-missing

If the report shows untested branches in your retry or parser logic, that is the gap to close.

Pitfalls and tradeoffs

• Over-mocking: If you fake too much, you test your mock, not the agent. Keep the fake’s schema identical to the real provider’s.
• Flaky eval confusion: Unit testing LLM agents is not the same as running evals. Unit tests check logic; evals check quality. Keep them in separate suites.
• Cost control: Never call a real API in unit tests. Use pytest markers to segregate integration tests that hit a live endpoint, and run those nightly.
• Cache-control: If you forward provider cache-control hints (like cache_control on system blocks), assert your agent sets them. A gateway that honors client routing directives will pass them through; your test should confirm you set them.

Putting it together

A minimal pytest file:

def test_agent_happy_path():
    fake = FakeLLM([tool_call_resp("get_weather", '{"city":"SF"}')])
    agent = Agent(fake)
    resp = agent.step("weather in SF?")
    assert agent.history[-1]["role"] == "assistant"
    assert fake.last_request["tools"] is not None

def test_agent_handles_bad_json():
    fake = FakeLLM([content_resp("{'city': 'SF'}")])  # invalid json
    agent = Agent(fake)
    with pytest.raises(ValueError):
        agent.step("weather?")

Unit testing LLM agents is mostly classical software testing applied to a fuzzy boundary. Draw the boundary sharply, mock at the edge, and assert on contracts. The model stays a black box; your code does not get a free pass.