Posted on
• Artificial Intelligence

Artificial Intelligence Testing with pytest

Author
  • User
    linuxbash
    Posts by this author
    Posts by this author

Artificial Intelligence Testing with pytest: Make Your Models Trustworthy on Linux

Ever shipped a model that scored 92% locally, only to see it dip to 88% in CI a week later? Or had a “flaky” test that passed yesterday and fails today—without any code changes? AI systems are uniquely sensitive to randomness, floating-point quirks, data drift, and performance regressions. The fix isn’t a new framework—it’s disciplined, automated testing with pytest.

In this guide, you’ll learn why AI testing is different, how pytest fits perfectly into a Linux-based workflow, and you’ll walk away with practical, reusable test patterns you can drop into your repos today.

Why test AI with pytest?

  • Reproducibility beats surprises: Seed control and deterministic pipelines stop “works on my machine” issues.

  • Guarantees over guesses: Tolerance-based assertions let you check correctness without demanding bit-for-bit equality.

  • Defend your SLAs: Benchmark inference and catch latency regressions before your users do.

  • Ship with confidence: Automated tests integrated into CI guard your accuracy and data contracts over time.

Quick install: system packages and Python tooling

You’ll want Python 3, venv, pip, a compiler toolchain, and then pytest plus a few helpful plugins.

Debian/Ubuntu (apt):

sudo apt update
sudo apt install -y python3 python3-venv python3-pip build-essential gcc

Fedora/RHEL/CentOS (dnf):

sudo dnf install -y python3 python3-pip python3-virtualenv gcc gcc-c++ make

openSUSE (zypper):

sudo zypper refresh
sudo zypper install -y python3 python3-pip python3-virtualenv gcc gcc-c++ make

Create and activate a virtual environment:

python3 -m venv .venv
. .venv/bin/activate

Install Python packages:

pip install --upgrade pip
pip install pytest pytest-cov pytest-benchmark pytest-xdist pytest-rerunfailures numpy scikit-learn

Optional, if you use pandas or PyTorch:

pip install pandas torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

A minimal, realistic AI example

Let’s keep the example small but real: an Iris classifier using scikit-learn with a pipeline for preprocessing + model.

Create ai_model/model.py:

from dataclasses import dataclass
from typing import Tuple

import numpy as np
from sklearn.datasets import load_iris
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

@dataclass
class TrainResult:
    model: object
    X_test: np.ndarray
    y_test: np.ndarray
    accuracy: float

def load_data() -> Tuple[np.ndarray, np.ndarray]:
    data = load_iris()
    return data.data.astype(np.float64), data.target.astype(np.int64)

def build_pipeline(random_state: int = 0):
    # StandardScaler + LogisticRegression is deterministic with liblinear + seed
    return make_pipeline(
        StandardScaler(),
        LogisticRegression(random_state=random_state, solver="liblinear", max_iter=1000)
    )

def train(random_state: int = 0) -> TrainResult:
    X, y = load_data()
    X_tr, X_te, y_tr, y_te = train_test_split(
        X, y, test_size=0.2, random_state=random_state, stratify=y
    )
    pipe = build_pipeline(random_state=random_state)
    pipe.fit(X_tr, y_tr)
    y_hat = pipe.predict(X_te)
    acc = accuracy_score(y_te, y_hat)
    return TrainResult(model=pipe, X_test=X_te, y_test=y_te, accuracy=acc)

def predict(model, X: np.ndarray) -> np.ndarray:
    # Return class probabilities to make tolerance-based checks meaningful
    return model.predict_proba(X)

5 actionable testing patterns that work for AI

1) Make randomness reproducible everywhere

Seed Python, NumPy, and your ML libraries inside pytest fixtures. Also disable multi-thread nondeterminism in CI via environment variables.

Create tests/conftest.py:

import os
import random
import numpy as np
import pytest

@pytest.fixture(autouse=True)
def seed_all():
    # Keep Python hashing stable across runs
    os.environ["PYTHONHASHSEED"] = "0"

    # Limit threaded math libraries for determinism (CI-friendly)
    os.environ.setdefault("OMP_NUM_THREADS", "1")
    os.environ.setdefault("MKL_NUM_THREADS", "1")
    os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
    os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")

    seed = 0
    random.seed(seed)
    np.random.seed(seed)

    try:
        import torch
        torch.manual_seed(seed)
        torch.use_deterministic_algorithms(True, warn_only=True)
    except Exception:
        pass

    yield

Bonus: In CI, add this to your job’s shell step:

export PYTHONHASHSEED=0 OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 NUMEXPR_NUM_THREADS=1

2) Test your data contracts and preprocessing

Catch schema and range issues early—before training silently degrades.

Create tests/test_data_contract.py:

import numpy as np
from ai_model.model import load_data, build_pipeline

def test_load_data_shapes_and_types():
    X, y = load_data()
    assert X.ndim == 2 and y.ndim == 1
    assert X.shape[0] == y.shape[0]
    assert X.dtype == np.float64
    assert y.dtype in (np.int32, np.int64)
    assert not np.isnan(X).any()

def test_pipeline_contains_scaler_and_classifier():
    pipe = build_pipeline()
    names = [name for name, _ in pipe.steps]
    assert "standardscaler" in names
    assert "logisticregression" in names

3) Use tolerance-based assertions, not exact equality

Floating-point math and randomized algorithms won’t be bit-exact. Use relative/absolute tolerances for probabilities, coefficients, and metrics.

Create tests/test_training_and_metrics.py:

import numpy as np
from ai_model.model import train, build_pipeline, load_data

def test_training_is_reproducible():
    # Train twice with the same seed and compare coefficients within tolerance
    r1 = train(random_state=0)
    r2 = train(random_state=0)

    # Extract coefficients from logistic regression inside the pipeline
    clf1 = r1.model.named_steps["logisticregression"]
    clf2 = r2.model.named_steps["logisticregression"]

    assert np.allclose(clf1.coef_, clf2.coef_, rtol=1e-6, atol=1e-8)
    assert np.allclose(clf1.intercept_, clf2.intercept_, rtol=1e-6, atol=1e-8)

def test_accuracy_threshold_guardrail():
    # If this trips, something changed: data, preprocessing, or model
    r = train(random_state=0)
    assert r.accuracy >= 0.90

4) Keep small “goldens” to catch semantic regressions

Golden outputs (small arrays, JSON blobs) catch subtle behavior changes. Store just a few rows to keep it light.

Create tests/test_golden_inference.py:

import os
import numpy as np
import pathlib
import pytest

from ai_model.model import train, predict

GOLDEN = pathlib.Path(__file__).parent / "goldens" / "probas_first5.npy"

@pytest.mark.filterwarnings("ignore::numpy.VisibleDeprecationWarning")
def test_inference_matches_golden_with_tolerance(tmp_path):
    r = train(random_state=0)
    probas = predict(r.model, r.X_test)

    # Focus on first 5 rows to keep goldens tiny
    first5 = probas[:5]

    # If golden missing, create it once intentionally (local dev convenience).
    # In CI, the golden should exist and be version-controlled.
    if not GOLDEN.exists():
        GOLDEN.parent.mkdir(parents=True, exist_ok=True)
        np.save(GOLDEN, first5)
        pytest.skip("Golden created; re-run tests to validate.")

    golden = np.load(GOLDEN)
    assert first5.shape == golden.shape
    # Tolerances account for minor numeric drift but catch meaningful changes
    assert np.allclose(first5, golden, rtol=1e-4, atol=1e-6)

Tip:

  • Commit tests/goldens/probas_first5.npy to your repo.

  • If you knowingly change behavior (e.g., new preprocessing), regenerate and review the diff.

5) Measure and lock down performance with pytest-benchmark

Latency is a feature. Track inference speed to spot slowdowns.

Create tests/test_performance.py:

import numpy as np
import pytest
from ai_model.model import train, predict

@pytest.mark.benchmark(group="inference")
def test_inference_latency_benchmark(benchmark):
    r = train(random_state=0)
    X = r.X_test

    def run():
        return predict(r.model, X)

    # The benchmark plugin stores timing; you can compare across commits with
    # --benchmark-compare or maintain a baseline. Avoid hard thresholds here.
    result = benchmark(run)
    assert result.shape[0] == X.shape[0]

Run locally:

pytest -q
pytest -q --benchmark-only
pytest -q --benchmark-compare

Parallelize tests and tame flakiness when needed:

pytest -q -n auto --reruns 2 --reruns-delay 1

Generate reports for CI:

pytest -q --cov=ai_model --cov-report=term-missing --junitxml=reports/junit.xml

Useful bash one-liners for CI

  • Fail fast if unseeded randomness sneaks in:
export PYTHONHASHSEED=0 OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 NUMEXPR_NUM_THREADS=1
pytest -q -n auto --maxfail=1
  • Cache dependencies and run benchmarks only on main:
if [ "$GITHUB_REF_NAME" = "main" ]; then
  pytest --benchmark-only --benchmark-save=main
else
  pytest --benchmark-only --benchmark-compare=main
fi

Real-world tips

  • Keep tests small: Use tiny datasets or sample production data. Avoid GPU in unit tests; keep that for integration/perf suites.

  • Mock external services: For LLM calls or HTTP-dependent steps, use pytest’s monkeypatch or libraries like responses to avoid network flakiness.

  • Separate test tiers: Fast unit tests on every commit, heavier benchmarks and integration tests on schedule or on main.

  • Version data and models: Pair your tests with data snapshots (DVC, git-lfs) for real reproducibility.

Conclusion and next steps

AI testing isn’t optional—your users will test in production if you don’t. With pytest on Linux, you can:

  • Make runs reproducible,

  • Assert correctness with sensible tolerances,

  • Guard your accuracy and latency over time,

  • And integrate it all cleanly into CI.

Your next step: 1) Drop the snippets above into your repo. 2) Wire up pytest -q --cov=... in your CI. 3) Add one golden test and one benchmark. 4) Watch flakiness and regressions disappear.

Have a tip or a gnarly edge case? Share it—and let’s make AI testing as routine (and boring) as any other software test.

Related posts