Should machine learning engineers write unit tests?

Yes, machine learning engineers should unit test feature transforms, data splits and serving code, because bugs there change predictions without an error. The trained model's quality needs an evaluation set instead, because each retraining changes its outputs.

How do you unit test a machine learning model?

Unit test the code around a machine learning model, not its predictions: feed hand-built rows through each transform and assert the exact output. Check the trained model by scoring a fixed evaluation set before each deployment.

What should I unit test in an ML pipeline?

In an ML pipeline, unit test each feature transform, the train and test split, the fitting of scalers on training rows only, and the parity of training and serving features. These tests run on five-row DataFrames without a GPU.

How do I catch training-serving skew?

Catch training-serving skew with a CI test that feeds one raw record through the training path and the serving path and compares the feature vectors. In production, compare the distributions of logged serving features with the training data, and alert when they diverge.

Should I write unit tests for machine learning code?

Verdict

Yes

Yes, unit test the feature transforms, data splits and serving code of a machine learning pipeline, including a parity test that feeds one raw record through the training and serving feature code.

Why

Blast radiususers
Change frequencyregularly
Detectabilityeventually
Reversibilitywith-effort
Test costmoderate

Yes, unit test the code around the model: feature transforms, data splits and serving code. For a model behind a product feature, Blast radius is users, and Change frequency is regularly, because features and training settings change about once a month. Detectability is eventually: a bug in feature code crashes nothing, the model still trains, and its score drops a few points that look like noise. Reversibility is with-effort, because a fix means retraining and redeploying. Test cost is moderate, about an hour for a module of transforms tested on five-row DataFrames, and rule R11 gives Test.

When the decision changes
When	Decision	Why
The check is whether the trained model predicts well enough: accuracy, ranking quality, drift	Test it differently: gate each deployment on a fixed evaluation set, and monitor live predictions	Test cost rises to heavy: each training run gives different weights, so a unit test cannot pin the model's output
The model approves loans, screens job applicants or flags patients for follow-up	Test mandatory: unit tests for every feature transform, evaluation for each group of people, and a second reviewer	Blast radius rises to safety-or-legal: a wrong feature can discriminate against people or harm a patient
The code lives in a research notebook that nobody deploys	Do not test; fix the random seed, check results by eye, and add tests when the code moves into the pipeline	Blast radius falls to none: only the author reads the output
A script recomputes features over historical data once, before a retraining	Test it differently: run the backfill on a copy and compare feature distributions with the old values	Change frequency falls to once: a unit test of the backfill would never run again

What breaks if you don't test

A scaler fit on the whole dataset leaks test rows into training, and the offline score goes up instead of down. The scikit-learn guide to common pitfalls shows feature selection with leaked test rows scoring 0.76 accuracy on random data, where chance gives 0.5. A category encoded differently in serving feeds the model inputs it never saw, the service keeps answering with status 200, and conversion drops for weeks before anyone looks.

What you lose if you over-test

A test that asserts the model's prediction for one input breaks on every retraining, and the team learns to paste in each new number unread. A test that trains a small model in CI slows every run and still passes when accuracy falls two points. Tests of whether StandardScaler computes a mean repeat the scikit-learn test suite.

How to test

The model is a small part of a production ML system, as Sculley and colleagues show in Hidden Technical Debt in Machine Learning Systems. The rest is data handling, tested with pytest on tiny hand-built inputs in CI:

Each feature transform: a normal row, a missing value, an unseen category.
The split: no record in both training and test data, no future row in training.
Fitting: scalers and encoders learn from training rows only.
Parity: one raw record gives the same vector in training and serving.

Before each deployment, score the new model on a fixed evaluation set and block the release when it scores below the current model. Google's ML Test Score rubric lists data, model and monitoring checks for production models.

Procedure and references

When the answer changes

The model's output touches money, credit, hiring or health: test every transform and add a second reviewer.
The code stays in a notebook that no product reads: add tests when it moves into the pipeline.

Real incident + Code example

The trailing space that cost five weeks

On a recommendation service I worked on, the training job lowercased and trimmed product categories before encoding them, and the serving code, written later by another developer, did not. A category sent as "Garden " fell into the unknown bucket at serving time. Dashboards stayed green, and clicks on the category carousel fell by about 9 percent for five weeks, until an analyst compared feature distributions in training data and serving logs. The fix was one shared encoding function and these tests:

import pandas as pd
from features import build_batch, build_online, split_by_date

def test_training_and_serving_build_the_same_vector():
    raw = {"category": "Garden ", "price": 12.5}
    batch = build_batch(pd.DataFrame([raw]))
    assert batch.iloc[0].tolist() == build_online(raw).tolist()

def test_unseen_category_goes_to_unknown_bucket():
    row = build_online({"category": "Toys", "price": 1.0})
    assert row["category_id"] == 0  # 0 is the unknown bucket

def test_no_future_rows_in_training():
    df = pd.DataFrame({"day": [1, 2, 3, 4], "clicks": [5, 7, 2, 9]})
    train, test = split_by_date(df, cutoff_day=3)
    assert train["day"].max() < 3 <= test["day"].min()

The parity test would have failed on the first serving commit.

Sources

Should I test data science code?Yes
Should I test prompts?Test it differently
Should AI write unit tests?Yes
Should I test helper functions?Yes

FAQ

Should machine learning engineers write unit tests?: Yes, machine learning engineers should unit test feature transforms, data splits and serving code, because bugs there change predictions without an error. The trained model's quality needs an evaluation set instead, because each retraining changes its outputs.
How do you unit test a machine learning model?: Unit test the code around a machine learning model, not its predictions: feed hand-built rows through each transform and assert the exact output. Check the trained model by scoring a fixed evaluation set before each deployment.
What should I unit test in an ML pipeline?: In an ML pipeline, unit test each feature transform, the train and test split, the fitting of scalers on training rows only, and the parity of training and serving features. These tests run on five-row DataFrames without a GPU.
How do I catch training-serving skew?: Catch training-serving skew with a CI test that feeds one raw record through the training path and the serving path and compares the feature vectors. In production, compare the distributions of logged serving features with the training data, and alert when they diverge.