🧭 How toSoftwareIntermediate✨ AI-assisted

How to Mock LLM API Calls in Pytest Without Flaky Tests

WittyTech Daily··5 min read
#pytest#testing#llm#python

Your CI went red because a language model phrased something differently, and nothing in your code changed. This guide shows how to mock LLM calls in pytest so unit tests run offline, cost nothing and return the same answer every time, including a tool call, a 429 and a timeout. You'll end with one thin client function, one fixture, a guard that fails any test that sneaks onto the network, and a small live suite for the nightly build.

The examples use Python and the Anthropic SDK, but the pattern is the same for any provider.

1. Put the LLM client behind one thin function

A test seam starts with a single function the rest of your code calls. Wrap the SDK once, return plain dataclasses, and translate provider errors into your own exception types. Your app code never touches SDK objects, so your fakes never have to build them.

# app/llm.py
import os
from dataclasses import dataclass, field
from functools import cache

import anthropic


class LLMRateLimited(Exception):
    """The provider returned a 429."""


class LLMTimeout(Exception):
    """The request timed out."""


@dataclass
class ToolCall:
    id: str
    name: str
    args: dict


@dataclass
class LLMReply:
    text: str = ""
    tool_calls: list[ToolCall] = field(default_factory=list)


@cache
def _client() -> anthropic.Anthropic:
    return anthropic.Anthropic(max_retries=0, timeout=20.0)


def complete(messages: list[dict], *, tools: list[dict] | None = None) -> LLMReply:
    kwargs = {"tools": tools} if tools else {}
    try:
        resp = _client().messages.create(
            model=os.environ["LLM_MODEL"],
            max_tokens=1024,
            messages=messages,
            **kwargs,
        )
    except anthropic.RateLimitError as exc:
        raise LLMRateLimited(str(exc)) from exc
    except anthropic.APITimeoutError as exc:
        raise LLMTimeout(str(exc)) from exc

    reply = LLMReply()
    for block in resp.content:
        if block.type == "text":
            reply.text += block.text
        elif block.type == "tool_use":
            reply.tool_calls.append(ToolCall(block.id, block.name, block.input))
    return reply

I set max_retries=0 because the SDK retries 429s and timeouts twice by default. Owning the retry policy in your own code keeps it visible to your tests. Set LLM_MODEL to whichever model ID you use.

Here is the code under test. It handles a tool call, a 429 and a timeout:

# app/support.py
from app import llm

BUSY_MESSAGE = "We're swamped right now. Try again in a minute."
TIMEOUT_MESSAGE = "That took too long. Please try again."

TOOLS = [{
    "name": "lookup_order",
    "description": "Look up an order by its ID.",
    "input_schema": {
        "type": "object",
        "properties": {"order_id": {"type": "string"}},
        "required": ["order_id"],
    },
}]


def lookup_order(order_id: str) -> str:
    return f"Order {order_id} shipped yesterday."  # stand-in for your real lookup


def answer(question: str) -> str:
    messages = [{"role": "user", "content": question}]
    for _ in range(2):  # one retry on timeout
        try:
            reply = llm.complete(messages, tools=TOOLS)
        except llm.LLMTimeout:
            continue
        except llm.LLMRateLimited:
            return BUSY_MESSAGE
        if reply.tool_calls and reply.tool_calls[0].name == "lookup_order":
            return lookup_order(**reply.tool_calls[0].args)
        return reply.text
    return TIMEOUT_MESSAGE

Note from app import llm and llm.complete(...). That choice matters in step 3.

2. Build a fake that plays back a script

A fake that always says "ok" can't test a 429. Make it replay whatever each test queues, exceptions included, and record every call so you can assert on what your code sent.

3. Mock the LLM in pytest with a fake client fixture

# tests/conftest.py
import pytest

from app import llm


class FakeLLM:
    """Plays back queued replies in order and records every call."""

    def __init__(self):
        self.script = []
        self.calls = []

    def queue(self, *items):
        self.script.extend(items)

    def __call__(self, messages, *, tools=None):
        self.calls.append({"messages": messages, "tools": tools})
        if not self.script:
            raise AssertionError("LLM called but no canned reply was queued")
        item = self.script.pop(0)
        if isinstance(item, Exception):
            raise item
        return item


@pytest.fixture
def fake_llm(monkeypatch):
    fake = FakeLLM()
    monkeypatch.setattr(llm, "complete", fake)
    return fake


@pytest.fixture(autouse=True)
def block_real_llm(request, monkeypatch):
    if request.node.get_closest_marker("live"):
        return

    def refuse(*args, **kwargs):
        raise AssertionError("Test hit the real LLM. Request the fake_llm fixture.")

    monkeypatch.setattr(llm, "complete", refuse)

The pytest monkeypatch docs promise that every change is undone after the test, so nothing leaks between tests. setattr also raises AttributeError if the name doesn't exist, so a typo can't silently patch nothing.

The autouse guard runs first and installs a function that fails loudly. A test that forgets fake_llm breaks in milliseconds instead of quietly spending money.

4. Test the happy path, a tool call, a 429 and a timeout

# tests/test_support.py
from app import llm, support
from app.llm import LLMReply, ToolCall


def test_tool_call_runs_the_order_lookup(fake_llm):
    fake_llm.queue(LLMReply(tool_calls=[
        ToolCall(id="toolu_1", name="lookup_order", args={"order_id": "A-100"})
    ]))
    assert "A-100" in support.answer("Where is order A-100?")
    assert fake_llm.calls[0]["tools"] == support.TOOLS


def test_429_returns_busy_message_without_retrying(fake_llm):
    fake_llm.queue(llm.LLMRateLimited("429"))
    assert support.answer("hi") == support.BUSY_MESSAGE
    assert len(fake_llm.calls) == 1


def test_one_timeout_is_retried(fake_llm):
    fake_llm.queue(llm.LLMTimeout("slow"), LLMReply(text="Hello!"))
    assert support.answer("hi") == "Hello!"
    assert len(fake_llm.calls) == 2


def test_two_timeouts_give_up_politely(fake_llm):
    fake_llm.queue(llm.LLMTimeout("slow"), llm.LLMTimeout("slow"))
    assert support.answer("hi") == support.TIMEOUT_MESSAGE

These are deterministic unit tests: no network, no API key, no tokens. Run pytest and the whole file finishes in well under a second.

5. Keep a small live suite out of the default run

Fakes prove your code handles replies. They can't prove the real API still returns what the fakes pretend. Cover that with a few live tests behind a marker that the default run skips.

# pyproject.toml
[tool.pytest.ini_options]
markers = ["live: calls the real LLM API, costs money, may be flaky"]
addopts = "-m 'not live'"
# tests/live/test_live_llm.py
import os

import pytest

from app import llm, support

pytestmark = [
    pytest.mark.live,
    pytest.mark.skipif(
        not (os.environ.get("ANTHROPIC_API_KEY") and os.environ.get("LLM_MODEL")),
        reason="needs ANTHROPIC_API_KEY and LLM_MODEL",
    ),
]


def test_model_requests_the_order_lookup_tool():
    reply = llm.complete(
        [{"role": "user", "content": "Where is order A-100?"}],
        tools=support.TOOLS,
    )
    assert reply.tool_calls and reply.tool_calls[0].name == "lookup_order"
pytest            # every push: offline, free, milliseconds
pytest -m live    # nightly job: real API, secrets injected by CI

The -m live on the command line overrides the one in addopts. The live test asserts on shape (which tool was called), never on wording.

Gotchas

  • Patch where the name is looked up. If support.py did from app.llm import complete, patching llm.complete changes nothing there. Patch app.support.complete instead, or keep the module-style import above. The unittest.mock docs explain where to patch.
  • Fakes can drift from the real signature. Mine copies complete's signature exactly. If you'd rather not maintain that, unittest.mock.create_autospec(llm.complete, side_effect=[...]) enforces the real signature and raises any exception you put in the list.
  • Copy canned replies from real responses. Print a real tool call once and match its shape, or your fakes will test a world that doesn't exist.
  • Never gate merges on the live suite. Run it nightly and read failures as a signal about the provider, not about your pull request.

Unit tests that never touch the network let you change retry and fallback logic with confidence, which is exactly the code that matters on the day a provider has a bad afternoon.

← More in Software