Testing ML code with pytest
Assert contracts and invariants. Mock the LLM. Gate quality with a floor.
Why it matters: You cannot assert exact LLM text. You can assert JSON shape, score ranges, and accuracy floors.
Basics
pytest basics
Name files test_*.py and functions test_*. Use plain assert.
Arrange inputs, act (call the function), assert the result.
One behaviour per test. Name the expectation in the function name.
Tiny example
Assert a score stays inside [0, 1].
def clamp(x):
return max(0.0, min(1.0, x))
assert 0.0 <= clamp(1.5) <= 1.0
print("ok", clamp(1.5))Basics
Fixtures and parametrize
Fixtures share setup across tests (sample data, temp folders).
@pytest.mark.parametrize runs one test on many inputs.
Cover edges: empty, whitespace, long text, non-ASCII.
Tiny example
Loop a few edge inputs the way parametrize would.
cases = ["", " ", "ok"]
for text in cases:
print(repr(text), "empty?", not text.strip())Basics
Testing non-deterministic AI
Assert what stays stable:
- Contract - JSON keys and types
- Invariants - chunking keeps all text; scores stay in range
- Thresholds - accuracy on a small golden set stays above a floor
Mock the LLM in unit tests. Keep live API checks separate.
Tiny example
Assert a JSON contract, not an exact string.
import json
raw = '{"label": "pass", "score": 0.91}'
data = json.loads(raw)
assert "label" in data and 0 <= data["score"] <= 1
print(data["label"])Then AI / ML
AI / ML examples
Same idea, used in real AI work. Press Try in lab to run it.
AI / ML example
AI / ML: mock the LLM client
Test parsing without network calls or token spend.
In AI / ML: Prompt builders and JSON parsers are tested against a fake client.
example.py
Python
import json
class FakeLLM:
def __init__(self, reply):
self.reply = reply
def complete(self, prompt):
return self.reply
def extract_entities(llm, text):
raw = llm.complete(f"entities for: {text}")
data = json.loads(raw)
return data["entities"]
print(extract_entities(FakeLLM('{"entities": ["Python"]}'), "Python ML"))AI / ML example 2
AI / ML: accuracy regression gate
Fail when a golden set drops below a floor.
In AI / ML: CI runs a tiny golden set and blocks deploy if accuracy falls.
example.py
Python
golden = [
{"q": "2 + 2", "expected": "4"},
{"q": "capital of France", "expected": "paris"},
]
def fake_system(q):
return {"2 + 2": "4", "capital of France": "Paris"}[q]
hits = sum(1 for row in golden if fake_system(row["q"]).lower() == row["expected"])
accuracy = hits / len(golden)
floor = 0.8
print("accuracy", accuracy)
print("PASS" if accuracy >= floor else "FAIL")Takeaways
- 1Test contracts, invariants, and thresholds - not exact LLM wording.
- 2Mock the model in unit tests; run live checks separately.
- 3Parametrize edge cases: empty, whitespace, huge, non-ASCII.