Unit testing AI agents with pytest forces you to confront the core problem: an agent is a loop that interleaves nondeterministic model outputs with stateful tool calls. If you let either touch the network, your tests become flaky and slow. The fix is to mock both the LLM client and the tools, then assert on the agent’s control flow as if it were any other state machine.
Step 1: Define a testable agent boundary
Most agent frameworks obscure the decision loop behind decorators and global state. Before you write a single test, refactor the agent so the LLM client and the tool registry are injected dependencies. This makes the agent a pure-ish function of (messages, tools) and lets you swap implementations in tests without monkey-patching imports.
A minimal agent is just a while loop that calls the model, checks for tool_calls, executes them, and appends results back to the message list. Keep the serialization logic in one place:
from typing import Callable, Dict, Any
class Agent:
def __init__(self, llm_client, tools: Dict[str, Callable]):
self.llm = llm_client
self.tools = tools
def run(self, user_msg: str) -> str:
messages = [{"role": "user", "content": user_msg}]
while True:
resp = self.llm.chat(messages)
msg = resp["choices"][0]["message"]
if msg.get("tool_calls"):
for call in msg["tool_calls"]:
fn = self.tools[call["function"]["name"]]
try:
result = fn(**call["function"]["arguments"])
except Exception as e:
result = f"error: {e}"
messages.append(msg)
messages.append({"role": "tool", "content": str(result)})
else:
return msg["content"]The try/except around the tool call is not optional—real tools fail, and your agent must serialize that failure back into the conversation. That branch is exactly what you will unit test later. If you use Pydantic models to validate tool arguments, do it inside the loop before calling fn; your tests can then confirm validation errors are also wrapped.
Step 2: Install pytest and create fixtures
You need pytest and a mocking helper. pytest-mock gives you the mocker fixture; respx is useful if you call an HTTP gateway directly. For pure in-process clients, unittest.mock is enough.
pip install pytest pytest-mock respxPut shared mock tools in conftest.py so every test can use them without repetition:
# conftest.py
import pytest
@pytest.fixture
def fake_tools():
def weather(city: str) -> str:
return f"Sunny in {city}"
return {"weather": weather}A good fixture returns a real Python callable, not a mock, for the happy path. Reserve mocks for cases where you need to assert on call arguments or simulate failures. This keeps the majority of tests readable and avoids over-specifying internals.
Step 3: Mock the LLM client with scripted responses
The LLM is the biggest source of nondeterminism. Replace it with a MagicMock whose chat method returns a queue of canned responses via side_effect. This turns the agent loop into a deterministic finite state machine that you can step through.
def test_agent_calls_tool(mocker, fake_tools):
llm = mocker.MagicMock()
llm.chat.side_effect = [
{"choices": [{"message": {"role": "assistant", "content": "",
"tool_calls": [{"function": {"name": "weather",
"arguments": {"city": "Berlin"}}}]}}]},
{"choices": [{"message": {"role": "assistant",
"content": "It's sunny in Berlin."}}]},
]
agent = Agent(llm, fake_tools)
out = agent.run("What's the weather?")
assert out == "It's sunny in Berlin."
assert llm.chat.call_count == 2If your production agent talks to an OpenAI-compatible gateway (n4n.ai exposes one endpoint covering 240+ models with automatic fallback), you would mock the HTTP requests.Session or the SDK client at the same boundary—never patch the agent internals. The assertion targets stay identical: correct number of calls, correct message shape.
Mocking a remote gateway over HTTP
If the agent posts to a URL directly, use respx to intercept:
import respx
import httpx
@respx.mock
def test_agent_via_http(fake_tools):
route = respx.post("https://api.example.com/v1/chat/completions").mock(
side_effect=[
httpx.Response(200, json={"choices": [{"message": {"role": "assistant",
"content": "", "tool_calls": [{"function": {"name": "weather",
"arguments": {"city": "Oslo"}}}]}}]}),
httpx.Response(200, json={"choices": [{"message": {"role": "assistant",
"content": "Cold in Oslo."}}]}),
]
)
client = httpx.Client()
agent = Agent(client, fake_tools) # assumes Agent uses client.post().json()
out = agent.run("weather Oslo")
assert out == "Cold in Oslo."
assert route.call_count == 2This exercises the real HTTP client code while keeping the model response fixed.
Step 4: Mock tools to verify side-effect capture
Tools are where agents cause real-world damage: sending emails, charging cards, mutating databases. In unit tests, replace the tool body with a mock that records its arguments and returns a fixed payload.
def test_tool_receives_correct_args(mocker):
mocked_weather = mocker.MagicMock(side_effect=lambda city: "Rain")
tools = {"weather": mocked_weather}
llm = mocker.MagicMock()
llm.chat.side_effect = [
{"choices": [{"message": {"role": "assistant", "content": "",
"tool_calls": [{"function": {"name": "weather",
"arguments": {"city": "Paris"}}}]}}]},
{"choices": [{"message": {"role": "assistant",
"content": "Rain in Paris."}}]},
]
agent = Agent(llm, tools)
agent.run("Weather in Paris?")
mocked_weather.assert_called_once_with(city="Paris")This test proves the agent parsed the model’s tool_calls correctly and forwarded the extracted arguments. If the model schema drifts, this fails before any integration environment sees it.
Step 5: Parameterize across decision paths
A single happy-path test is weak. Use pytest.mark.parametrize to exercise multiple LLM scripts and tool outputs without duplicating boilerplate.
import pytest
@pytest.mark.parametrize("tool_out,expected", [
("Sunny", "Sunny in Tokyo"),
("Snow", "Snow in Tokyo"),
("Typhoon", "Typhoon in Tokyo"),
])
def test_agent_parametrized(mocker, tool_out, expected):
tools = {"weather": lambda city: tool_out}
llm = mocker.MagicMock()
llm.chat.side_effect = [
{"choices": [{"message": {"role": "assistant", "content": "",
"tool_calls": [{"function": {"name": "weather",
"arguments": {"city": "Tokyo"}}}]}}]},
{"choices": [{"message": {"role": "assistant",
"content": expected}}]},
]
agent = Agent(llm, tools)
assert agent.run(f"weather {tool_out}") == expectedParameterization keeps your suite honest: if the agent suddenly ignores the tool result and hardcodes a string, all three cases catch it.
Step 6: Assert on error handling and prompt construction
The agent’s error branch is the most important code you own—the LLM vendor and the tool vendor are black boxes. Force a tool exception and confirm the agent wraps it and continues the loop.
def test_agent_surfaces_tool_error(mocker):
def boom(city): raise RuntimeError("API down")
tools = {"weather": boom}
llm = mocker.MagicMock()
llm.chat.side_effect = [
{"choices": [{"message": {"role": "assistant", "content": "",
"tool_calls": [{"function": {"name": "weather",
"arguments": {"city": "Rome"}}}]}}]},
{"choices": [{"message": {"role": "assistant",
"content": "Could not get weather."}}]},
]
agent = Agent(llm, tools)
out = agent.run("weather Rome?")
assert out == "Could not get weather."
# verify the error string was injected as a tool message
assert "error: API down" in llm.chat.call_args_list[1].args[0][-1]["content"]Here we inspect the second chat call’s messages to confirm the agent serialized the exception into the tool role. That is the contract between your code and the model. Add a similar assertion that the system prompt or prior user messages are preserved across turns.
Step 7: Run the suite and verify determinism
Execute the tests in quiet mode and confirm zero network activity. If you have a linter or coverage tool, wire it in, but the baseline is green.
pytest -qExpected output:
....... [100%]
7 passed in 0.09sIf you see 7 passed in under a second and no requests log lines, your mocks are isolating the agent correctly. Any test that hits a real endpoint is an integration test and should live in a separate directory marked # integration.
What good coverage looks like
After following these steps, your suite should include: happy path with tool, tool argument correctness, parameterized tool outputs, tool exception handling, and at least one test asserting the exact message list sent to the LLM. That last one catches prompt regressions that silently degrade agent quality. Unit testing AI agents with pytest is not about proving the model is smart; it is about proving your glue code routes data correctly when the model behaves as specified.
When you later add a real gateway call in staging, keep these mocks as the gatekeeper. Only promote a change to integration once the mocked suite is green. That discipline is the difference between shipping an agent that occasionally works and one whose control flow you actually understand.