pytest — Basics and Fixtures
Tests, fixtures, parametrization, and the patterns that keep an ML codebase testable.
TL;DR
pytest is the de facto Python testing framework. A test is a
function whose name starts with test_ and which uses plain assert
statements. There’s no class hierarchy to inherit from, no
setUp/tearDown boilerplate, no assertEquals. Just functions and
asserts.
The four features that make it the choice:
assertreports the actual values when it fails. No need forassertEqual.- Fixtures for reusable setup. Dependency-injected by parameter name.
parametrizefor the same test against many inputs.- A massive plugin ecosystem.
pytest-asynciofor async,pytest-mockfor mocks,pytest-covfor coverage,pytest-xdistfor parallelism.
The minimum viable test file:
# tests/test_features.py
import pytest
from my_project.features import normalize
def test_normalize_zero_centers():
out = normalize([1.0, 2.0, 3.0])
assert sum(out) == pytest.approx(0.0)
def test_normalize_handles_empty():
with pytest.raises(ValueError):
normalize([])
Run: pytest. Done. No config, no setup module.
The picture in your head
pytest discovers tests by walking the project, looking for files
named test_*.py (or *_test.py), inside which it picks up
def test_* functions and class Test* classes. Each test runs in
isolation; if it raises, it failed; if it doesn’t, it passed.
project/
├── src/my_project/...
└── tests/
├── conftest.py # shared fixtures
├── test_features.py # tests
└── test_train.py
conftest.py is special: pytest auto-imports any fixtures you define
there into every test in the same directory and below. No imports
needed in test files.
Assertions
assert is enough. Pytest rewrites the bytecode so failures show the
values of the expressions, not just “assertion failed”:
def test_addition():
assert 1 + 1 == 3
# Output:
# def test_addition():
# > assert 1 + 1 == 3
# E assert 2 == 3
Useful pytest extensions:
import pytest
# Approximate float equality
assert sum(values) == pytest.approx(0.0, abs=1e-6)
# Expect an exception
with pytest.raises(ValueError, match="must be positive"):
sqrt(-1)
# Expect a warning
with pytest.warns(DeprecationWarning):
old_function()
Fixtures — reusable setup
A fixture is a function decorated with @pytest.fixture that returns
something a test needs. Tests get fixtures by naming them as
arguments — pytest’s dependency injection.
# tests/conftest.py
import pytest
import torch
@pytest.fixture
def small_model():
return torch.nn.Linear(4, 2)
@pytest.fixture
def random_batch():
torch.manual_seed(0)
return torch.randn(8, 4)
# tests/test_train.py
def test_forward_shape(small_model, random_batch):
out = small_model(random_batch)
assert out.shape == (8, 2)
small_model and random_batch are constructed fresh per test (default
scope). pytest traces dependencies — if a fixture depends on another
fixture, it gets injected too.
Fixture scopes
| Scope | Constructed once per | Use for |
|---|---|---|
function (default) | Each test | Cheap setup; default. |
class | Each test class | Setup shared across tests in a class. |
module | Each test file | Expensive setup (load a model, open a DB). |
session | Whole pytest run | Very expensive (download a dataset). |
@pytest.fixture(scope="session")
def trained_model():
# Loaded once for the entire test run
return load_pretrained("test-model.pt")
@pytest.fixture(scope="module")
def database():
db = make_test_db()
yield db # tests run with db
db.cleanup() # teardown after the module's tests finish
The yield form gives you setup and teardown — code before yield
runs at start, code after runs after the last test that used it.
parametrize — the same test, many inputs
@pytest.mark.parametrize("input,expected", [
([1, 2, 3], 6),
([], 0),
([-1, 1], 0),
([0.5], 0.5),
])
def test_sum(input, expected):
assert sum(input) == expected
Pytest generates four separate tests, each named with the parameters, each shown individually in the output. If three pass and one fails, you see exactly which.
For multiple parametrizations stacked:
@pytest.mark.parametrize("x", [1, 2, 3])
@pytest.mark.parametrize("y", [10, 20])
def test_pair(x, y):
assert x * y > 0
# generates 6 tests: (x=1,y=10), (x=1,y=20), (x=2,y=10), ...
For ML, parametrize is invaluable for “test this loss against several
input shapes,” “test this tokenizer against several languages,” “test
this model interface against several backends.”
Mocking — pytest-mock (or stdlib unittest.mock)
def test_calls_api(mocker): # 'mocker' is the pytest-mock fixture
mock_client = mocker.patch("my_project.llm.openai_client")
mock_client.complete.return_value = "fake response"
result = my_function("prompt")
assert result == "fake response"
mock_client.complete.assert_called_once_with("prompt")
mocker.patch(...) replaces the named attribute for the duration of
the test, then restores it. The test runs without hitting the real API.
For ML specifically: mock the LLM client, mock the database, mock the
filesystem (or use tmp_path — a pytest builtin). Don’t mock NumPy
or PyTorch — those are deterministic and fast; just call them.
Useful built-in fixtures
| Fixture | What it gives you |
|---|---|
tmp_path | A unique pathlib.Path to a fresh temp dir for the test. Auto-cleaned. |
tmp_path_factory | Session-scoped temp-dir factory. |
monkeypatch | Set/unset env vars and module attrs; restored after the test. |
capsys / capfd | Capture stdout / stderr. |
caplog | Capture log records. |
def test_writes_checkpoint(tmp_path):
out = tmp_path / "model.pt"
save_checkpoint(out)
assert out.exists()
assert out.stat().st_size > 0
def test_uses_env(monkeypatch):
monkeypatch.setenv("MODEL_NAME", "test-model")
assert load_settings().model_name == "test-model"
Markers — categorising tests
@pytest.mark.slow
def test_expensive_thing(): ...
@pytest.mark.skip(reason="not implemented yet")
def test_future(): ...
@pytest.mark.skipif(sys.platform == "win32", reason="linux only")
def test_linux_only(): ...
@pytest.mark.xfail(reason="known bug, fixing in #123")
def test_known_broken(): ...
Run only specific markers: pytest -m "slow" or pytest -m "not slow".
For ML, mark expensive (GPU, full-data) tests with @pytest.mark.slow
and exclude them from the fast PR-checking suite; run them in a
nightly job.
A worked example — testing an ML preprocessing module
# src/my_project/preprocessing.py
import numpy as np
import pandas as pd
def normalize_features(df: pd.DataFrame, cols: list[str]) -> pd.DataFrame:
if df.empty:
raise ValueError("df is empty")
out = df.copy()
for c in cols:
mean = df[c].mean()
std = df[c].std()
if std == 0:
raise ValueError(f"column {c} has zero variance")
out[c] = (df[c] - mean) / std
return out
# tests/test_preprocessing.py
import numpy as np
import pandas as pd
import pytest
from my_project.preprocessing import normalize_features
@pytest.fixture
def sample_df():
return pd.DataFrame({
"a": [1.0, 2.0, 3.0, 4.0],
"b": [10.0, 20.0, 30.0, 40.0],
"c": ["x", "y", "z", "w"],
})
def test_normalizes_to_zero_mean(sample_df):
out = normalize_features(sample_df, ["a", "b"])
assert out["a"].mean() == pytest.approx(0.0)
assert out["b"].mean() == pytest.approx(0.0)
def test_preserves_other_columns(sample_df):
out = normalize_features(sample_df, ["a"])
assert (out["c"] == sample_df["c"]).all()
assert (out["b"] == sample_df["b"]).all()
def test_does_not_mutate_input(sample_df):
original = sample_df.copy()
normalize_features(sample_df, ["a"])
pd.testing.assert_frame_equal(sample_df, original)
def test_raises_on_empty_df():
with pytest.raises(ValueError, match="empty"):
normalize_features(pd.DataFrame(), ["a"])
def test_raises_on_zero_variance():
df = pd.DataFrame({"x": [1.0, 1.0, 1.0]})
with pytest.raises(ValueError, match="zero variance"):
normalize_features(df, ["x"])
@pytest.mark.parametrize("n_rows", [1, 10, 100, 1000])
def test_works_for_various_sizes(n_rows):
df = pd.DataFrame({"x": np.random.default_rng(0).standard_normal(n_rows)})
if n_rows > 1:
out = normalize_features(df, ["x"])
assert out["x"].std() == pytest.approx(1.0, abs=1e-6)
Notice: tests are short, each verifies one thing, fixtures are reused, edge cases (empty, zero variance) are explicit, parameterised tests cover sizes. This is what a healthy ML test file looks like.
Common gotchas
unittest.TestCaseand pytest fixtures don’t mix well. Stick to plain functions for new code.conftest.pyper directory. Fixtures in aconftest.pyare available only in that directory and below. Put shared ones attests/conftest.py.- Fixture not used. Pytest silently doesn’t run a fixture you didn’t request. Make sure the test function arg name matches the fixture name exactly.
monkeypatchin module-scope fixtures.monkeypatchis function-scoped only. Usemockerfrom pytest-mock, or build a custom session-scoped patch.- Network or GPU in tests. Slows the suite to uselessness. Mock
network. Mark GPU tests
@pytest.mark.slowand exclude from PR runs. - Tests that depend on each other. A failure in one cascading through to fake-fail another is impossible to debug. Each test is independent.
assertfor runtime checks in production code.assertis removed bypython -O. Useif not cond: raise ValueError(...).
Where pytest shows up in real ML codebases
pytest tests/in CI, on every PR.pytest-asynciofor testing async LLM batch logic.pytest-covfor coverage reports; not a goal but a useful artifact.pytest-xdistfor parallel test execution (-n autoruns on all cores).pytest --pdbdrops you into the debugger at the point of failure — invaluable when a test only fails in CI.hypothesisfor property-based tests (great for ML invariants: “softmax of any input sums to 1”, “zip then unzip = identity”).
The defensive habit: when you fix a bug, write the test that would have caught it first. The test that fails before the fix and passes after is the one that prevents regressions.
Resources
- pytest documentation — docs.pytest.org — canonical, very good.
- pytest fixtures reference — docs.pytest.org — full fixture mechanism.
- Brian Okken — Python Testing with pytest (2nd ed.) — pragprog.com — best book on pytest.
- pytest-mock — pypi.org/project/pytest-mock —
mockerfixture wrappingunittest.mock. - hypothesis — hypothesis.readthedocs.io — property-based testing.
- pytest-xdist — pytest-xdist.readthedocs.io — parallel test execution.