Write your first tests for an LLM app

Beginner → Intermediate · 15 min · tested with pytest 9

"How do you test something that gives a different answer every time?" You mostly don't test the model. You test your code around it: the prompt you build, how you parse the reply, and what happens when things go wrong. That code is ordinary Python, and it's where most bugs live. These tests are fast, free and run on every commit.

pip install pytest

---

1. Make the model swappable

The trick that makes LLM code testable: don't create the API client deep inside your functions. Pass in "anything that can complete a prompt". In production that's the real model; in tests it's a fake.

support.py

import json
from dataclasses import dataclass

CATEGORIES = {"refund", "delivery", "complaint", "other"}

def build_prompt(question: str, chunks: list[str], max_chunks: int = 3) -> str:
    """Build the RAG prompt: numbered context, then the question."""
    context = "\n".join(f"[{i}] {c}" for i, c in enumerate(chunks[:max_chunks], start=1))
    return f"Answer using only the context. Cite sources like [1].\n\nContext:\n{context}\n\nQuestion: {question}"

def parse_category(reply: str) -> str:
    """Read the model's JSON reply; fall back to 'other' on anything unexpected."""
    try:
        category = json.loads(reply).get("category", "other")
    except (json.JSONDecodeError, AttributeError):
        return "other"
    return category if category in CATEGORIES else "other"

@dataclass
class SupportBot:
    llm: object                      # anything with .complete(prompt) -> str

    def classify(self, message: str) -> str:
        reply = self.llm.complete(f'Classify into {sorted(CATEGORIES)}. Reply as JSON {{"category": ...}}.\n\n{message}')
        return parse_category(reply)

    def answer(self, question: str, chunks: list[str]) -> str:
        if not chunks:
            return "Sorry, I couldn't find that in our help articles."
        try:
            return self.llm.complete(build_prompt(question, chunks))
        except TimeoutError:
            return "The assistant is busy right now. Please try again in a minute."

---

2. A fake LLM as a fixture

A fixture is something pytest builds and hands to any test that asks for it by name. Put shared fixtures in tests/conftest.py:

tests/conftest.py

import pytest

class FakeLLM:
    """Stands in for the real model: returns scripted replies and records prompts."""
    def __init__(self, replies=None, error=None):
        self.replies = list(replies or [])
        self.error = error
        self.prompts = []

    def complete(self, prompt: str) -> str:
        self.prompts.append(prompt)
        if self.error:
            raise self.error
        return self.replies.pop(0)

@pytest.fixture
def fake_llm():
    return FakeLLM

The fake records the prompts it received, so tests can check what your code sent, not just what came back.

---

3. Write the tests

tests/test_support.py

import os
import pytest
from support import SupportBot, build_prompt, parse_category

# --- plain functions: the easiest, most valuable tests ---

def test_build_prompt_numbers_chunks_and_keeps_question():
    prompt = build_prompt("How long do refunds take?", ["Refunds take 5 days.", "Delivery takes 3 days."])
    assert "[1] Refunds take 5 days." in prompt
    assert "[2] Delivery takes 3 days." in prompt
    assert prompt.endswith("Question: How long do refunds take?")

def test_build_prompt_limits_chunks():
    prompt = build_prompt("q", [f"chunk {i}" for i in range(10)], max_chunks=3)
    assert "[3] chunk 2" in prompt and "[4]" not in prompt

# --- one test, many cases: every odd reply a model has ever sent you goes here ---

@pytest.mark.parametrize("reply, expected", [
    ('{"category": "refund"}', "refund"),
    ('{"category": "Refund"}', "other"),
    ('{"category": "cancel"}', "other"),
    ("Sure! It's a refund.", "other"),
    ("[]", "other"),
])
def test_parse_category(reply, expected):
    assert parse_category(reply) == expected

# --- the bot, with the fake model ---

def test_classify_uses_model_reply(fake_llm):
    llm = fake_llm(replies=['{"category": "delivery"}'])
    assert SupportBot(llm).classify("Where is my parcel?") == "delivery"
    assert "Where is my parcel?" in llm.prompts[0]

def test_answer_without_context_never_calls_model(fake_llm):
    llm = fake_llm()
    reply = SupportBot(llm).answer("What is the warranty?", chunks=[])
    assert "couldn't find" in reply
    assert llm.prompts == []                  # no wasted API call

def test_answer_handles_timeout(fake_llm):
    bot = SupportBot(fake_llm(error=TimeoutError()))
    assert "try again" in bot.answer("Refund time?", ["Refunds take 5 days."])

# --- a real-model test: only runs when a key is available ---

@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="needs OPENAI_API_KEY")
def test_real_model_classifies_refund():
    from openai import OpenAI

    class OpenAILLM:
        def complete(self, prompt):
            r = OpenAI().chat.completions.create(
                model="gpt-4o-mini", temperature=0, response_format={"type": "json_object"},
                messages=[{"role": "user", "content": prompt}])
            return r.choices[0].message.content

    assert SupportBot(OpenAILLM()).classify("My order arrived broken, I want my money back") == "refund"

---

4. Run them

python -m pytest -q          # "python -m" puts the project folder on the import path

..........s                                                              [100%]
10 passed, 1 skipped in 0.03s

Ten tests in 0.03 seconds, no API key, no cost. The s is the real-model test, skipped because no key is set. See why with -rs:

python -m pytest -v -rs

tests/test_support.py::test_parse_category[{"category": "refund"}-refund] PASSED [ 27%]
tests/test_support.py::test_parse_category[{"category": "Refund"}-other] PASSED [ 36%]
...
tests/test_support.py::test_real_model_classifies_refund SKIPPED (ne...) [100%]
SKIPPED [1] tests\test_support.py:40: needs OPENAI_API_KEY

What a failure looks like

Suppose someone "simplifies" parse_category to return category.lower(). The tests catch it immediately:

>       assert parse_category(reply) == expected
E       AssertionError: assert 'cancel' == 'other'

FAILED tests/test_support.py::test_parse_category[{"category": "Refund"}-other]
FAILED tests/test_support.py::test_parse_category[{"category": "cancel"}-other]
2 failed, 3 passed in 0.05s

The model invented a category (cancel) and the new code would have passed it straight to your routing logic.

---

5. What to test in an LLM app

Test	Speed and cost	Runs
Prompt building, parsing, validation, chunking, routing	Milliseconds, free	Every commit
Your error handling (timeouts, bad JSON, empty context) with a fake model	Milliseconds, free	Every commit
A handful of real-model smoke tests (temperature=0)	Seconds, a few paise	On PRs, when a key is set
A full eval suite (accuracy on 50-500 examples)	Minutes, costs money	Nightly or before releases

Real-model tests should check properties, not exact wording: the category is refund, the answer cites [1], the JSON validates. Every time a weird model reply causes a bug in production, add it to a parametrize list so it can never break you again.

---

Checklist

• The LLM client is passed in, not created inside functions.
• Prompt building and reply parsing have plain unit tests.
• A fake model covers timeouts, invalid replies and empty context.
• Real-model tests skip cleanly without a key and check properties, not exact text.
• pytest runs in CI on every pull request.

More in Testing with pytest and Evaluating agents.

Next: Debug Python in VS Code →