Learn AI Create the Future — with Pranay
Main site

Testing ML Code and CI

Unit testing machine learning code

A unit test checks one small function on its own, before you trust it inside a bigger pipeline — the same habit that catches a mis-set kitchen scale before it ruins a whole dish.

Read these first

On this page 7
  1. The short answer
  2. The analogy you have already lived
  3. Why it exists
  4. How it works
  5. A real example you have seen
  6. Remember this
  7. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

The short answer

Unit testing means checking one small piece of code on its own, before you trust it inside anything bigger.

The analogy you have already lived

You have used a kitchen scale. Before weighing rice for a recipe, you press the button that sets the display back to zero with nothing on it. If it does not read zero, you know before you cook, not after the dish is too salty.

A unit test is that same zero-check, written down so a computer can run it every time. It runs on every small function in your code, not only once, by hand, when you remember.

Why it exists

A machine learning pipeline is built from many small functions. One cleans a column. One clips outlier values. One turns a raw record into the row shape the model expects.

Each function is small enough to get right. The problem is that a pipeline chains ten or twenty of them, and a mistake in any single one does not crash anything. It quietly changes the numbers instead.

You do not find out from an error message. You find out three weeks later, when accuracy has dropped and nobody can say which of the twenty functions is the one at fault. A unit test checks each function the moment it is written, so a broken one is caught in seconds, by name.

How it works

   your pipeline
   ---------------------------------------------
   raw record -> clip_income -> years_to_bucket -> to_feature_row -> model
   ---------------------------------------------
        ^              ^               ^
        |              |               |
     test this    test this       test this
     on its own   on its own      on its own,
                                   separately from
                                   the model

Each arrow gets its own small, fast check. None of them needs the model, the database, or the network to run.

A real example you have seen

Photo apps resize an image before a model looks at it. If the resize function ever swaps width and height by mistake, every photo trained or scored after that point is silently distorted.

A one-line unit test — "resizing a 200x100 image gives back 100x50, not 50x100" — catches that the moment the function is written. It does not wait until after a model has been trained on thousands of sideways photos.

Remember this

  • A unit test checks one small function, in isolation, without needing the whole pipeline running.
  • It exists because pipeline bugs are usually silent — they change numbers, they do not crash anything.
  • Cheap, fast checks written once catch mistakes that would otherwise take weeks to trace back to their source.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install pytest pandas

The functions being tested

Three small, single-purpose helpers from a loan-approval pipeline. Each does exactly one job, which is what makes it easy to test on its own.

features.py
"""Small feature-engineering helpers for a loan-approval model."""
import pandas as pd


def clip_income(income: float, lo: float = 0.0, hi: float = 200.0) -> float:
    """Caps a reported monthly income (thousands) to a sane range."""
    return max(lo, min(hi, income))


def years_to_bucket(years: float) -> str:
    """Turns raw years of credit history into one of three bands."""
    if years < 2:
        return "new"
    if years < 7:
        return "established"
    return "veteran"


def to_feature_row(applicant: dict) -> pd.DataFrame:
    """Builds the one-row frame the model expects, in the model's column order."""
    row = {
        "income": clip_income(applicant["income"]),
        "years": applicant["years"],
        "age": applicant["age"],
    }
    return pd.DataFrame([row], columns=["income", "years", "age"])

The tests

test_features.py
from features import clip_income, years_to_bucket, to_feature_row


def test_clip_income_leaves_normal_values_alone():
    assert clip_income(60.0) == 60.0


def test_clip_income_caps_a_typo_of_a_million():
    # a real data-entry bug: someone enters rupees, not thousands of rupees
    assert clip_income(1_000_000.0) == 200.0


def test_clip_income_floors_a_negative_value():
    assert clip_income(-40.0) == 0.0


def test_years_to_bucket_boundaries():
    assert years_to_bucket(0.0) == "new"
    assert years_to_bucket(1.99) == "new"
    assert years_to_bucket(2.0) == "established"
    assert years_to_bucket(6.99) == "established"
    assert years_to_bucket(7.0) == "veteran"


def test_to_feature_row_has_the_right_shape_and_order():
    row = to_feature_row({"income": 60.0, "years": 7.0, "age": 34})
    assert list(row.columns) == ["income", "years", "age"]
    assert row.shape == (1, 3)


def test_to_feature_row_applies_the_income_clip():
    row = to_feature_row({"income": -5.0, "years": 7.0, "age": 34})
    assert row.loc[0, "income"] == 0.0
bash
pytest test_features.py -q
Output
......                                                                   [100%]
6 passed in 0.55s

That exact timing is one run on one machine — yours will differ by a few tenths of a second, but the pass count will not.

Line-by-line: the boundary tests

test_years_to_bucket_boundaries checks 1.99, 2.0, 6.99 and 7.0 — the exact numbers next to each if. Bugs hide at boundaries far more often than in the middle of a range. A test on 5.0 would pass even if someone changed < 7 to < 6 by accident. A test at 6.99 and 7.0 would not.

Now break something on purpose

Swap max and min inside clip_income — an easy typo that still runs without error:

python
# in features.py, clip_income becomes:
return min(lo, max(hi, income))
Output
FF....                                                                   [100%]
================================== FAILURES ===================================
_________________ test_clip_income_leaves_normal_values_alone _________________

    def test_clip_income_leaves_normal_values_alone():
>       assert clip_income(60.0) == 60.0
E       assert 0.0 == 60.0
E        +  where 0.0 = clip_income(60.0)

test_features.py:5: AssertionError
__________________ test_clip_income_caps_a_typo_of_a_million __________________

    def test_clip_income_caps_a_typo_of_a_million():
>       assert clip_income(1_000_000.0) == 200.0
E       assert 0.0 == 200.0
E        +  where 0.0 = clip_income(1000000.0)

test_features.py:10: AssertionError
=========================== short test summary info ===========================
FAILED test_features.py::test_clip_income_leaves_normal_values_alone - assert...
FAILED test_features.py::test_clip_income_caps_a_typo_of_a_million - assert 0...
2 failed, 4 passed in 0.78s

The swapped function now returns 0.0 for every input, no matter what you pass it. Nothing crashed. Nothing printed a warning. Without these two tests, every income in the pipeline would silently become zero, and the model trained on that data would look plausible while being completely wrong.

Common mistakes

Testing the whole pipeline and calling it a unit test. If a test needs the trained model, a database, or a network call, it is an integration test, not a unit test. Keep unit tests small enough to run in milliseconds, so nobody is tempted to skip them.

Only testing the happy path. clip_income(60.0) passing tells you almost nothing. The value of a test is concentrated at the edges: negative numbers, zero, huge numbers, boundary values.

One giant test function. test_everything() that checks five unrelated things fails with one vague message. Five small tests fail with five exact ones. pytest -q above shows exactly that — the failure names alone tell you which function broke.

Testing implementation instead of behaviour. Do not assert that clip_income "calls min then max". Assert what it returns for a given input. That way the test still protects you if the internals are rewritten to do the same job differently.

Try it yourself

Add a function is_valid_age(age) that rejects ages under 18 or over 100, and write three tests for it: a normal age, and both boundaries. Then break the boundary on purpose, the way clip_income was broken above, and watch which test catches it.

What to learn next

Researcher — Mathematics and papers.

What a unit test is, formally

A unit test exercises a function $f$ under a chosen input $x$ and asserts a property of $f(x)$ — usually equality with a known-correct value, but sometimes a weaker property such as a range or an invariant. The unit is whatever is small enough to reason about without a mental model of the rest of the system; in practice this is a single function or method, not a module.

Coverage as a lower bound, not a target

Line coverage is the fraction of source lines executed by the test suite. It is a floor, not a goal: 100% line coverage with no assertions on the results (assert result is not None) proves the code ran, not that it is correct. Treat a coverage number as evidence of untested code (the uncovered lines) rather than evidence of correct code.

Test data generation

Property-based testing, implemented for Python by Hypothesis (MacIver, 2019), generates many inputs from a declared strategy and searches for a counterexample, rather than requiring the author to enumerate cases by hand:

python
from hypothesis import given, strategies as st

@given(st.floats(min_value=-1e9, max_value=1e9, allow_nan=False))
def test_clip_income_is_always_in_range(income):
    assert 0.0 <= clip_income(income) <= 200.0

This single test subsumes the three hand-written clip_income boundary tests above, and additionally searches for a counterexample across the full float range. The trade-off is speed: property tests run many iterations per test, so they belong in a slower CI stage than the hand-written suite.

Where unit testing sits in the ML Test Score

Breck et al. (2017), The ML Test Score, places conventional unit testing under the Model Development axis, but explicitly separates it from the harder ML-specific axes: data tests, model-quality tests, and infrastructure tests. Unit testing alone — no matter how thorough — cannot detect a training-serving skew, a label flip, or a drifted feature distribution. Those require the behavioural and data tests covered in the lessons that follow this one.

Papers

What to learn next

Related terms

What to learn next

These follow on from what you just read.

  • Testing ML Code and CI

    Testing a training loop

    A training-loop test runs the whole loop for a few fast steps on fake data, so a broken loss function or a crashed checkpoint is caught in seconds instead of after a multi-hour run.

  • Testing ML Code and CI

    Behavioural tests for models

    A behavioural test checks what a trained model does on cases you choose on purpose, the way a driving examiner tests specific manoeuvres instead of only checking the car starts.

  • Testing ML Code and CI

    Golden outputs and regression tests

    A golden file freezes today's known-good output so tomorrow's "small refactor" cannot quietly change it — the same reason a tailor keeps your exact measurements on paper instead of guessing again each visit.