Skip to content

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:

  1. assert reports the actual values when it fails. No need for assertEqual.
  2. Fixtures for reusable setup. Dependency-injected by parameter name.
  3. parametrize for the same test against many inputs.
  4. A massive plugin ecosystem. pytest-asyncio for async, pytest-mock for mocks, pytest-cov for coverage, pytest-xdist for 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

ScopeConstructed once perUse for
function (default)Each testCheap setup; default.
classEach test classSetup shared across tests in a class.
moduleEach test fileExpensive setup (load a model, open a DB).
sessionWhole pytest runVery 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

FixtureWhat it gives you
tmp_pathA unique pathlib.Path to a fresh temp dir for the test. Auto-cleaned.
tmp_path_factorySession-scoped temp-dir factory.
monkeypatchSet/unset env vars and module attrs; restored after the test.
capsys / capfdCapture stdout / stderr.
caplogCapture 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.TestCase and pytest fixtures don’t mix well. Stick to plain functions for new code.
  • conftest.py per directory. Fixtures in a conftest.py are available only in that directory and below. Put shared ones at tests/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.
  • monkeypatch in module-scope fixtures. monkeypatch is function-scoped only. Use mocker from 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.slow and 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.
  • assert for runtime checks in production code. assert is removed by python -O. Use if not cond: raise ValueError(...).

Where pytest shows up in real ML codebases

  • pytest tests/ in CI, on every PR.
  • pytest-asyncio for testing async LLM batch logic.
  • pytest-cov for coverage reports; not a goal but a useful artifact.
  • pytest-xdist for parallel test execution (-n auto runs on all cores).
  • pytest --pdb drops you into the debugger at the point of failure — invaluable when a test only fails in CI.
  • hypothesis for 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