Testing Machine Learning Code and Data

Ordinary software tests check that code does what it was written to do. Machine learning systems need more, because two of their most important inputs are not code: the data, which changes without anyone editing it, and the trained model, whose behaviour nobody wrote down line by line. Google's ML Test Score paper makes the point directly: testing is key to production readiness, yet a model's behaviour is hard to specify in advance. It sets out 28 specific tests and monitoring needs drawn from production experience (Breck et al.: The ML Test Score).

This guide writes nine practical tests for a real model and runs them.

Three kinds of test

Nine tests, one real dataset

The model predicts daily bike rentals from the UCI Bike Sharing data, released under CC BY 4.0 (UCI: Bike Sharing). The tests use pytest fixtures to load the data and split it once per module, a pattern pytest documents for sharing expensive setup while keeping results repeatable (pytest: fixtures).

import numpy as np
import pandas as pd
import pytest
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error

FEATURES = ["season", "yr", "mnth", "holiday", "weekday", "workingday", "weathersit", "temp", "atemp", "hum", "windspeed"]


@pytest.fixture(scope="module")
def day():
    return pd.read_csv("data/day.csv")                          # UCI Bike Sharing (CC BY 4.0)


@pytest.fixture(scope="module")
def split(day):
    return day[day.dteday < "2012-07-01"], day[day.dteday >= "2012-07-01"]


def train(data, seed=42):
    return RandomForestRegressor(n_estimators=100, random_state=seed).fit(data[FEATURES], data.cnt)


# ---- data tests: is the input what the model expects? ----
def test_schema(day):
    assert set(FEATURES + ["dteday", "cnt"]) <= set(day.columns)
    assert not day[FEATURES + ["cnt"]].isna().any().any()


def test_ranges(day):
    assert day[["temp", "atemp", "hum", "windspeed"]].apply(lambda c: c.between(0, 1).all()).all()
    assert day.season.isin([1, 2, 3, 4]).all() and (day.cnt >= 0).all()


def test_counts_add_up(day):
    assert (day.casual + day.registered == day.cnt).all()


def test_one_row_per_day(day):
    dates = pd.to_datetime(day.dteday)
    assert dates.is_unique and (dates.diff().dropna().dt.days == 1).all()


def test_humidity_is_plausible(day):
    bad = day.loc[day.hum <= 0, "dteday"].tolist()
    assert not bad, f"humidity of 0% on {bad}: a sensor fault, not weather"


# ---- model tests: does the trained model behave sensibly? ----
def test_training_is_deterministic(split):
    train_df, test_df = split
    assert np.array_equal(train(train_df).predict(test_df[FEATURES]), train(train_df).predict(test_df[FEATURES]))


def test_beats_naive_baseline(split):
    train_df, test_df = split
    model_mae = mean_absolute_error(test_df.cnt, train(train_df).predict(test_df[FEATURES]))
    baseline_mae = mean_absolute_error(test_df.cnt, np.full(len(test_df), train_df.cnt.tail(30).mean()))
    assert model_mae < 0.8 * baseline_mae, f"model {model_mae:.0f} vs baseline {baseline_mae:.0f}"


def test_rain_lowers_predictions(split):
    train_df, test_df = split
    model = train(train_df)
    clear, rainy = test_df[FEATURES].assign(weathersit=1), test_df[FEATURES].assign(weathersit=3)
    assert model.predict(rainy).mean() < model.predict(clear).mean()


def test_predictions_in_range(split):
    train_df, test_df = split
    predictions = train(train_df).predict(test_df[FEATURES])
    assert (predictions >= 0).all() and predictions.max() <= 2 * train_df.cnt.max()


class Summary:
    # a tiny pytest plugin: collect each test's outcome and its failure message
    def __init__(self):
        self.rows = []

    def pytest_runtest_logreport(self, report):
        if report.when == "call":
            message = "" if report.passed else report.longrepr.reprcrash.message.splitlines()[0]
            self.rows.append((report.nodeid.split("::")[-1], report.outcome, message))


if __name__ == "__main__":
    summary = Summary()
    pytest.main([__file__, "-p", "no:terminal", "-p", "no:cacheprovider"], plugins=[summary])
    for name, outcome, message in summary.rows:
        print(f"{outcome.upper():<7} {name:<32} {message}")
    print(f"\n{sum(o == 'passed' for _, o, _ in summary.rows)} passed, {sum(o == 'failed' for _, o, _ in summary.rows)} failed")

PASSED  test_schema                      
PASSED  test_ranges                      
PASSED  test_counts_add_up               
PASSED  test_one_row_per_day             
FAILED  test_humidity_is_plausible       AssertionError: humidity of 0% on ['2011-03-10']: a sensor fault, not weather
PASSED  test_training_is_deterministic   
PASSED  test_beats_naive_baseline        
PASSED  test_rain_lowers_predictions     
PASSED  test_predictions_in_range        

8 passed, 1 failed

What the output shows

• The data passed the obvious checks. Every expected column is present, nothing is missing, weather values lie between 0 and 1, there is exactly one row per day, and casual plus registered riders always equals the total.
• One test caught a real fault. The humidity reading for 10 March 2011 is 0%. It is inside the valid range, so the range test passed it. It is also physically implausible for a weather station, so the plausibility test failed. This is exactly the kind of quiet error that slips into training data and skews a model.
• The model behaved as expected. Training was deterministic with a fixed seed, the model beat a naive baseline by the required margin, and its predictions fell when the weather was set to rain on every test day.
• Not every extreme is an error. The dataset also contains a day with only 22 rentals, 29 October 2012, recorded in bad weather. A test that rejected it would delete real behaviour. Data tests should target impossible values, not rare ones.

When a data test fails

A failing data test is a decision point, not just a red line. For each check, decide in advance what happens:

• Stop the pipeline when the data is unusable, such as a missing column.
• Quarantine the bad rows and continue, logging how many were removed.
• Repair them by a documented rule. For the humidity fault, treating 0% as missing and filling it from the neighbouring days is a reasonable policy.
• Alert a person when the cause needs investigating at the source.

Write the choice down next to the test, so the next person knows why the pipeline behaved as it did.

Behavioural tests

The rain test is a behavioural test. It does not check accuracy. It checks that the model responds in the direction common sense expects. The idea comes from software testing and was popularised for models by CheckList, whose authors found that held-out accuracy often overestimates performance; in their user study, practitioners using the method created twice as many tests and found almost three times as many bugs (Ribeiro et al.: CheckList).

Useful patterns:

• Directional: a change in one input should move the prediction one way, as rain lowers rentals.
• Invariance: a change that should not matter, such as an ID column, should leave predictions unchanged.
• Minimum functionality: simple, obvious cases must be right.

Where tests run

• Data tests at every ingestion, before training starts, so bad data stops the pipeline.
• Training and model tests in CI on every change to code or configuration. See CI/CD for machine learning.
• Quality gates before promotion, comparing the new model with the current one. See the model registry.
• Monitoring after release, because production data drifts after every test has passed. See monitoring ML models in production.

Choosing thresholds

A quality gate needs a threshold, and a threshold needs a reason. The baseline test above requires the model to be at least 20% better than predicting the recent average. Thresholds like that should come from the business cost of errors and from how much results vary between runs. A threshold tighter than the run-to-run spread will fail at random. See reproducibility in ML.

Common mistakes

• Testing only code and assuming the data is fine.
• Range checks without plausibility checks.
• Tests that remove real extremes along with errors.
• Accuracy as the only model test.
• Gates with arbitrary thresholds that fail at random.
• Tests that never run because they are not in CI.

Testing, CI/CD and automated deployment for ML are covered in the CI/CD for ML module of the MLOps course.

Sources

• The ML Test Score: A Rubric for ML Production Readiness and Technical Debt Reduction (Google Research) — Breck et al., Google
• Beyond Accuracy: Behavioral Testing of NLP models with CheckList (arXiv) — Ribeiro et al.
• How to use fixtures (pytest documentation) — pytest
• Bike Sharing dataset (UCI Machine Learning Repository) — UCI Machine Learning Repository