Parametrized Unit Tests and Reusable Pytest Fixtures

Tests are executable behavioral contracts

A unit test checks one small, isolated unit of behavior: usually a function, method, or class. It should answer a concrete question such as:

• Does channel normalization apply the correct statistic to each RGB channel?
• Does the function preserve the batch-image-channel shape?
• Does invalid input fail clearly rather than silently producing invalid values?

For the vectorized work from the previous lesson, a useful contract is:

Given channels-last images of shape (N, H, W, 3), RGB means, and nonzero RGB standard deviations, normalize_rgb returns an array of the same shape whose channels are normalized independently.

Suppose the implementation lives in src/ai_utils/normalization.py:

import numpy as np


def normalize_rgb(
    images: np.ndarray,
    mean: np.ndarray,
    std: np.ndarray,
) -> np.ndarray:
    if images.ndim != 4 or images.shape[-1] != 3:
        raise ValueError(
            "images must have shape (batch, height, width, 3)"
        )

    if mean.shape != (3,) or std.shape != (3,):
        raise ValueError(
            "mean and std must each have shape (3,)"
        )

    if np.any(std == 0):
        raise ValueError("std must contain only nonzero values")

    return (images - mean) / std

A conventional project layout keeps production code and tests separate:

project/
├── src/
│   └── ai_utils/
│       └── normalization.py
└── tests/
    └── test_normalization.py

pytest discovers files named test_*.py and functions named test_*. From the project root, run:

python -m pytest

Using python -m pytest helps ensure that pytest runs under the same Python environment as your project.

A strong unit test has three conceptual stages:

1. Arrange the input and dependencies.
2. Act by calling the unit under test.
3. Assert observable behavior.

For numerical AI code, assertions should usually cover both values and shapes. A correct-looking tensor with an unintended shape can easily cause a later transformer or vision pipeline failure.

Parametrization: one assertion pattern, many independently reported cases

It is tempting to write several nearly identical tests:

def test_normalize_mean_pixel() -> None:
    ...

def test_normalize_bright_pixel() -> None:
    ...

That is valid, but the repeated structure obscures the fact that these are instances of the same rule. pytest.mark.parametrize separates the test logic from its case data.

Here is a focused parametrized test for normalize_rgb:

import numpy as np
import pytest

from ai_utils.normalization import normalize_rgb


@pytest.mark.parametrize(
    ("pixel", "expected"),
    [
        pytest.param(
            [0.5, 0.25, 0.0],
            [0.0, 0.0, 0.0],
            id="channel-means-map-to-zero",
        ),
        pytest.param(
            [1.0, 0.75, 1.0],
            [1.0, 2.0, 1.0],
            id="bright-pixel-is-normalized-per-channel",
        ),
    ],
)
def test_normalize_rgb_pixel_values(
    rgb_stats: tuple[np.ndarray, np.ndarray],
    pixel: list[float],
    expected: list[float],
) -> None:
    mean, std = rgb_stats

    images = np.asarray(
        pixel,
        dtype=np.float32,
    )[None, None, None, :]

    actual = normalize_rgb(images, mean, std)

    expected_array = np.asarray(
        expected,
        dtype=np.float32,
    )[None, None, None, :]

    assert actual.shape == (1, 1, 1, 3)
    np.testing.assert_allclose(
        actual,
        expected_array,
        rtol=1e-6,
        atol=1e-6,
    )

The decorator creates two independent test items. pytest will report them separately, using the readable IDs:

test_normalize_rgb_pixel_values[channel-means-map-to-zero]
test_normalize_rgb_pixel_values[bright-pixel-is-normalized-per-channel]

This matters when a regression occurs: the failure output identifies the behavior that broke, rather than merely saying that a large generic test failed.

The parameter names in the decorator must match arguments in the test function:

@pytest.mark.parametrize(("pixel", "expected"), [...])
def test_normalize_rgb_pixel_values(pixel, expected):
    ...

Think of every case as a row in a data table:

Keep success and failure behavior separate

A frequent parametrization mistake is adding a boolean such as should_raise and then branching inside the test:

# Avoid this pattern.
if should_raise:
    ...
else:
    ...

The test has now acquired its own control flow, which makes its intent harder to inspect. Separate the success contract from the error contract instead:

@pytest.mark.parametrize(
    "bad_std",
    [
        pytest.param(
            [0.5, 0.0, 1.0],
            id="zero-green-standard-deviation",
        ),
        pytest.param(
            [0.0, 0.25, 1.0],
            id="zero-red-standard-deviation",
        ),
    ],
)
def test_normalize_rgb_rejects_zero_std(
    rgb_stats: tuple[np.ndarray, np.ndarray],
    bad_std: list[float],
) -> None:
    mean, _ = rgb_stats
    images = np.ones((1, 1, 1, 3), dtype=np.float32)

    with pytest.raises(ValueError, match="std"):
        normalize_rgb(
            images,
            mean,
            np.asarray(bad_std, dtype=np.float32),
        )

This test has one clear claim:

Every standard-deviation vector containing zero is invalid.

pytest.raises verifies that the failure is intentional and of the right type. match="std" additionally checks that the message communicates the relevant problem. The string is treated as a regular expression, so keep such matches specific but not unnecessarily brittle.

For numerical tests, favor:

np.testing.assert_allclose(actual, expected)

over:

assert actual == expected

Array equality does not produce a single useful Boolean result, and floating-point computations often differ by very small rounding errors even when mathematically equivalent.

Fixtures: named, reusable setup dependencies

The test above requests rgb_stats, but we have not defined it yet. A fixture is a function decorated with @pytest.fixture that creates a dependency for tests.

Place broadly reusable fixtures in tests/conftest.py:

import numpy as np
import pytest


@pytest.fixture
def rgb_stats() -> tuple[np.ndarray, np.ndarray]:
    mean = np.array([0.5, 0.25, 0.0], dtype=np.float32)
    std = np.array([0.5, 0.25, 1.0], dtype=np.float32)
    return mean, std

pytest sees this test signature:

def test_normalize_rgb_pixel_values(
    rgb_stats: tuple[np.ndarray, np.ndarray],
    pixel: list[float],
    expected: list[float],
) -> None:

and resolves its arguments in two different ways:

• pixel and expected come from @pytest.mark.parametrize.
• rgb_stats comes from the fixture named rgb_stats.

The fixture name acts as a dependency declaration. You do not call the fixture yourself with rgb_stats(). pytest calls it at the appropriate time and injects its returned value.

Why not use a global array?

This is fragile:

RGB_MEAN = np.array([0.5, 0.25, 0.0], dtype=np.float32)
RGB_STD = np.array([0.5, 0.25, 1.0], dtype=np.float32)

If any test mutates one of these arrays, that mutation may affect a later test. Then test results depend on execution order, which is a serious reliability problem.

By default, pytest fixtures have function scope. That means pytest creates a fresh fixture value for each test function. In the example, each test gets fresh mean and std arrays.

This isolation is especially important for mutable AI-engineering objects:

• NumPy arrays that a preprocessing function might modify in place,
• temporary datasets or cache directories,
• PyTorch models whose parameters or training mode can change,
• tokenizer settings,
• environment variables,
• mock service clients.

A useful design rule is:

A fixture should provide one clearly named dependency in a known state.

Good fixture names describe the provided object:

rgb_stats
small_token_batch
temporary_checkpoint_dir
inference_client

Avoid vague names such as:

setup
data
thing

They force readers to inspect implementation details before understanding the test.

Fixture dependencies and scope

Fixtures can request other fixtures. For example, pytest supplies the built-in tmp_path fixture, which gives each test a unique temporary directory. You can build a project-specific fixture on top of it:

from pathlib import Path

import pytest


@pytest.fixture
def experiment_dir(tmp_path: Path) -> Path:
    directory = tmp_path / "experiment"
    directory.mkdir()
    return directory

A test can now write files without knowing where temporary storage comes from:

def test_checkpoint_path_is_created(
    experiment_dir: Path,
) -> None:
    checkpoint_path = experiment_dir / "model.pt"
    checkpoint_path.write_bytes(b"checkpoint")

    assert checkpoint_path.exists()

pytest fixture scopes control how long a fixture value lives:

For most unit-test fixtures, use the default function scope. It is the safest choice.

A wider scope can reduce repeated expensive setup:

@pytest.fixture(scope="module")
def vocabulary() -> dict[str, int]:
    return {
        "<pad>": 0,
        "<bos>": 1,
        "image": 2,
        "caption": 3,
    }

Use broader scope only when the returned object is effectively read-only. A session-scoped GPU model that one test changes with .train(), weight updates, or mutable caches can contaminate every later test. For unit tests, it is often better to test model-adjacent logic with tiny deterministic tensors rather than loading a real large model.

Cleanup with yield fixtures

When a fixture owns a resource that must be released, use yield. Code before yield performs setup; code after it performs teardown, even if the test fails.

from collections.abc import Iterator
from pathlib import Path
from typing import TextIO

import pytest


@pytest.fixture
def opened_log(tmp_path: Path) -> Iterator[TextIO]:
    log_path = tmp_path / "events.log"

    with log_path.open("w+", encoding="utf-8") as handle:
        yield handle

This continues the resource-lifetime discipline from the context-manager lesson:

• The test receives an open file handle.
• Once the test is finished, the with block closes it.
• The built-in tmp_path fixture isolates the directory from other tests.

For databases, sockets, spawned processes, GPU allocations, or temporary environment changes, cleanup behavior is not optional infrastructure—it is part of a reliable test design.

Parametrizing fixtures versus parametrizing tests

There are two related but distinct tools:

1. Test parametrization varies a specific test’s cases.

   @pytest.mark.parametrize("temperature", [0.0, 0.7, 1.0])
   def test_sampling_accepts_temperature(
       temperature: float,
   ) -> None:
       ...
   

2. Fixture parametrization reruns every test that depends on the fixture for each fixture configuration.

   @pytest.fixture(
       params=["channels_last", "channels_first"],
   )
   def image_layout(request: pytest.FixtureRequest) -> str:
       return request.param
   

Fixture parametrization is appropriate when the fixture represents a configuration that many tests must support, such as a storage backend or supported image layout. But it expands the test matrix quickly: three parametrized test cases combined with two fixture parameters produce six test items.

Use ordinary test parametrization by default. Introduce fixture parametrization when an entire group of tests genuinely needs to run across the same configurations.

A practical testing workflow for numerical code

When adding or changing a vectorized function, build a compact behavioral suite rather than trying to test every possible array.

A useful sequence is:

1. Choose a small hand-computable example.
   The expected values should be obvious enough to verify without reusing the implementation’s logic.

2. Test a semantic edge case.
   For normalization, a channel equal to its mean should map to zero.

3. Test a structural invariant.
   For element-wise normalization, output shape should equal input shape.

4. Test invalid inputs deliberately.
   Zero standard deviation, wrong channel count, or incompatible statistic shape should produce clear failures.

5. Keep shared setup in fixtures.
   Centralize stable, reusable data such as RGB statistics, temporary paths, and miniature vocabularies.

6. Use parameter IDs for meaningful cases.
   A failure named zero-green-standard-deviation is easier to act on than one named [1].

7. Run targeted tests while developing.

   python -m pytest tests/test_normalization.py -q
   

   To inspect exactly what pytest will run:

   python -m pytest tests/test_normalization.py --collect-only -q
   

Avoid tests that simply duplicate the implementation’s internal steps. For example, computing expected normalization by calling the same helper used in production does not independently verify much. Prefer manually derived expected values for small cases, plus shape and error-contract checks.

Key takeaways

• A unit test is an executable statement of a small behavioral contract.
• @pytest.mark.parametrize turns one test pattern into multiple independently reported cases.
• Use descriptive pytest.param(..., id="...") case names to make failures understandable.
• Keep successful behavior and expected-error behavior in separate tests instead of branching inside a test.
• Fixtures provide named dependencies through test-function arguments; pytest resolves and injects them.
• Function-scoped fixtures provide fresh state for each test and prevent mutation leakage.
• Put reusable fixtures in conftest.py; use yield fixtures when setup owns a resource that needs cleanup.
• For numerical code, test values with np.testing.assert_allclose, along with shapes and invalid-input contracts.

Next, you will profile a Python program to identify its dominant CPU-time or memory bottleneck. The tests you wrote here will become a safety net: after optimizing code, you can verify that performance improvements did not alter numerical behavior.