Advanced20 min

Testing ML code with pytest

Write tests that catch data and logic regressions, including the non-deterministic parts of ML systems.

Why this matters in AI / ML / GenAI

You cannot assert that an LLM returns an exact string. You can assert the JSON parses, the schema matches, chunking never loses text, and accuracy stays above a floor. Teams that test their data and glue code ship far fewer incidents.

pytest basics

Put tests in tests/, name files test_*.py and functions test_*. Assert with plain assert — pytest rewrites it to show both sides on failure.

Run with pytest -q. Add --cov for coverage once the basics pass.

Arrange, Act, Assert: build inputs, call the function, check the result. One behaviour per test, and a name that states the expectation: test_chunker_preserves_all_text.

Fixtures and parametrize

@pytest.fixture builds shared setup (a sample DataFrame, a temp directory, a fake client) and injects it by parameter name.

@pytest.mark.parametrize runs one test over many inputs — perfect for edge cases: empty string, whitespace only, very long text, non-ASCII.

tmp_path is a built-in fixture giving a clean temp directory per test, so file tests never collide.

Testing non-deterministic AI systems

Four things you can reliably assert:

  1. Contract — output parses as JSON and has the required keys and types.
  2. Invariants — chunking loses no characters; embeddings have the expected dimension; scores stay in [0, 1].
  3. Thresholds — accuracy on a small golden set stays above a floor (a regression gate, not a unit test).
  4. Determinism where you control it — seed random number generators; set temperature=0 for reproducible generations.

Mock the LLM in unit tests. Real API calls make tests slow, flaky, and expensive. Test your prompt-building and response-parsing logic against a fake client, and run a small live suite separately on a schedule.

Copy-paste examples

Copy into your own editor, or load one into the compiler below and press Run.

Test file layout

Copy into tests/test_chunking.py and run pytest -q locally.

# tests/test_chunking.py
import pytest
from myproject.chunking import chunk_text

def test_chunker_preserves_all_text():
    text = "a" * 100
    chunks = chunk_text(text, size=30, overlap=0)
    assert "".join(chunks) == text

def test_chunker_respects_size():
    chunks = chunk_text("x" * 95, size=30, overlap=0)
    assert all(len(c) <= 30 for c in chunks)

@pytest.mark.parametrize("bad", ["", "   ", None])
def test_chunker_rejects_empty(bad):
    with pytest.raises((ValueError, TypeError)):
        chunk_text(bad, size=10)

Mock the LLM client

Test your parsing logic without spending tokens or waiting on the network.

import json

class FakeLLM:
    def __init__(self, reply):
        self.reply = reply
        self.calls = []

    def complete(self, prompt):
        self.calls.append(prompt)
        return self.reply

def extract_entities(llm, text):
    raw = llm.complete(f"Return JSON with key entities for: {text}")
    data = json.loads(raw)
    if "entities" not in data:
        raise ValueError("missing entities key")
    return data["entities"]

llm = FakeLLM('{"entities": ["Python", "Kubernetes"]}')
print(extract_entities(llm, "We run Python on Kubernetes"))
print("prompt sent:", llm.calls[0][:40], "...")

bad_llm = FakeLLM('{"stuff": []}')
try:
    extract_entities(bad_llm, "text")
except ValueError as err:
    print("caught:", err)

Assert an invariant, not an exact string

Runs here as plain asserts; the same body works inside pytest.

def chunk_text(text, size=30, overlap=0):
    if not text or not text.strip():
        raise ValueError("text must not be empty")
    step = size - overlap
    return [text[i:i + size] for i in range(0, len(text), step)]

text = "Testing ML glue code prevents silent data corruption."
chunks = chunk_text(text, size=20, overlap=0)

assert "".join(chunks) == text, "chunker lost characters"
assert all(len(c) <= 20 for c in chunks), "chunk too large"
assert len(chunks) == 3, f"expected 3 chunks, got {len(chunks)}"
print("all invariants hold:", chunks)

Accuracy regression gate

Run against a small golden set in CI; fail the build if quality drops.

golden = [
    {"q": "capital of France", "expected": "paris"},
    {"q": "2 + 2", "expected": "4"},
    {"q": "language of ML", "expected": "python"},
]

def fake_system(question):
    answers = {"capital of France": "Paris", "2 + 2": "4", "language of ML": "Java"}
    return answers[question]

hits = sum(1 for row in golden if fake_system(row["q"]).lower() == row["expected"])
accuracy = hits / len(golden)
floor = 0.8

print(f"accuracy={accuracy:.2f} floor={floor}")
if accuracy < floor:
    print("FAIL: quality regression — block the deploy")
else:
    print("PASS")

Write assertions for a scoring function

Try it — in-browser Python

Break normalize_score (return raw) and see which assertion fires first.

Output

Python runs in your browser. First run downloads the runtime.

Press Run (or Ctrl+Enter) to execute.

CPython in WebAssembly. Stdlib works. NumPy and pandas load on demand. No input(), no GPU, no network installs.

Takeaways

  • Test contracts, invariants, and thresholds — not exact LLM wording.
  • Mock the model in unit tests; run live checks separately.
  • parametrize covers edge cases: empty, whitespace, huge, non-ASCII.