Introduction

Unit testing machine learning code is less about proving that a model is "good" and more about proving that the software around the model behaves correctly. Good ML tests focus on deterministic pieces such as preprocessing, feature construction, data contracts, metric calculations, and model-serving glue rather than treating training outcomes themselves as ordinary unit-test assertions.

Test the Deterministic Parts First

Most ML systems have many components that should behave deterministically:

• feature scaling,
• tokenization,
• label mapping,
• train-test splitting helpers,
• metric calculations,
• prediction post-processing.

Those are excellent unit-test targets because a fixed input should give a fixed output.

1def normalize_age(age: int) -> float:
2    return age / 100.0
3
4
5def test_normalize_age():
6    assert normalize_age(25) == 0.25

That may look simple, but bugs in these small transformations often cause more production pain than model math itself.

Use Small, Fixed Fixtures

ML code should not require huge datasets in unit tests. Keep fixtures tiny and explicit so failures are easy to understand.

1import pandas as pd
2
3
4def build_feature_frame(df: pd.DataFrame) -> pd.DataFrame:
5    out = df.copy()
6    out["name_len"] = out["name"].str.len()
7    return out[["name_len"]]
8
9
10def test_build_feature_frame():
11    df = pd.DataFrame({"name": ["Ada", "Linus"]})
12    result = build_feature_frame(df)
13    assert list(result["name_len"]) == [3, 5]

Small fixtures make it clear what broke and avoid turning unit tests into slow integration jobs.

Control Randomness

If a test touches randomized behavior, fix the random seed or mock the randomness source.

1import numpy as np
2
3
4def sample_threshold():
5    rng = np.random.default_rng(123)
6    return rng.uniform()
7
8
9def test_sample_threshold():
10    assert round(sample_threshold(), 6) == 0.682352

Without deterministic seeds, test failures may be noisy and misleading.

Do Not Unit Test "Model Quality" Naively

A unit test should not normally say "training accuracy must always be 0.93" on a real stochastic pipeline. That kind of check is fragile and usually belongs in a higher-level validation suite with tolerances, controlled data snapshots, and separate execution expectations.

What you can test instead:

• that the trainer returns an object of the right type,
• that prediction shapes are correct,
• that probabilities sum correctly,
• that serialization and deserialization preserve behavior.

1import numpy as np
2
3
4def assert_binary_probs(probs: np.ndarray):
5    assert probs.shape[1] == 2
6    assert np.allclose(probs.sum(axis=1), 1.0)

Test Saved-Model Contracts

If you serve models in production, test the boundary around loading and predicting.

1class FakeModel:
2    def predict(self, features):
3        return [1 for _ in features]
4
5
6def serve(model, features):
7    preds = model.predict(features)
8    return {"predictions": preds}
9
10
11def test_serve():
12    result = serve(FakeModel(), [[0.1], [0.2]])
13    assert result == {"predictions": [1, 1]}

This kind of test protects deployment logic even when the real model object is large or expensive. It also keeps unit tests fast, which matters because slow tests are the ones teams quietly stop running.

Common Pitfalls

• Writing brittle tests against stochastic training outcomes.
• Using huge real datasets in unit tests instead of tiny fixtures.
• Failing to seed randomness in code paths that need deterministic assertions.
• Ignoring preprocessing and feature-engineering functions, which are often the most testable parts.
• Treating slow end-to-end training checks as unit tests instead of a separate validation layer.

Summary

• Unit-test deterministic ML code first: preprocessing, features, metrics, and serving glue.
• Use tiny fixed fixtures so failures are easy to diagnose.
• Control randomness with seeds or mocks.
• Avoid fragile unit tests that assert exact model quality numbers after training.
• Keep heavy model-validation checks separate from fast software-level unit tests.