Testing ML code
Vivek inherits the churn repo at an NBFC in Mangaluru. A teammate tunes a feature and, while editing, changes one divisor in spend_share. Nothing crashes. The API answers, the training script prints a score, the demo goes well. Three weeks later the recall review shows the model missing churners it used to catch. Git blame finds the edit in minutes; nothing had complained for twenty-one days. His lead's verdict: our code cannot tell us when it is wrong. Write tests that can.
- Regular code fails with a loud error, while ML code can fail as a wrong number that looks fine. Companies pay for people who can build the alarm that catches the second kind.
- Interviewers ask how you would test an ML system, and most freshers have no answer. After today you name four layers and open a repo where each one works.
- Module 4.6 wires this suite into CI, where it runs on every push. A gate that blocks a worse model from shipping is one of the most professional artifacts a fresher can show.
Most bugs you have met so far announced themselves. A missing file raised an error, a bad payload earned a 422, a syntax mistake refused to run at all. ML code has a second failure mode: the code runs, returns a number, and the number is wrong. Vivek's divisor bug crashed nothing, because 0.0025 is a perfectly legal float. A wrong number is worse than a crash, because a crash stops the line while a wrong number ships.
pytest is the tool that turns this silence into noise. A unit test — a tiny program that calls one of your functions with known input and checks the output against an answer you worked out by hand. An assertion — a line starting with assert that states what must be true; when it is false, the run stops and shouts. pytest collects files named test_*.py, runs the functions inside that start with test_, and prints a dot per pass and an F per failure. You will write your first one in the watch section, then break it on purpose to learn how failure reads.
Here is the contrast that decides whether changes stay safe. The wrong way is the eyeball check: edit the code, run the script, stare at the score, decide it feels fine. It costs attention every time, it misses small drops, and on leave days it does not happen at all. The right way is a test suite — the collection of all tests, run as one command — that fails loudly the moment any promise breaks. Machines are better than tired eyes at noticing that 0.824 became 0.816.
ML testing has four layers, and each answers a different question. Data validation tests ask: is the incoming data even legal — expected columns, sane ranges, no nulls, no duplicates? Feature unit tests ask: does each feature function still produce the numbers we worked out by hand? Pipeline-contract tests ask: does the sealed pipeline fit, predict the right shape and range, and survive a save and load round trip? The model-quality gate asks the question everyone forgets: is the model we are about to ship still good enough?
On top sits one smoke test — a quick end-to-end check that the whole service answers at all. It calls the FastAPI app from module 4.2 through TestClient, a fake caller that needs no running server. The layers matter because each catches what the others cannot. The smoke test cannot see a worse model, and the quality gate cannot see a broken route. When a planted bug turns exactly one layer red, the red tells you where to look.
Tests need data, and real data is the wrong data for most of them. A fixture — a small piece of known setup, like twelve hand-typed rows, that pytest hands to any test that asks for it. Tiny fixtures keep tests fast, and hand-typed values mean every expected answer can be checked on paper. The quality gate is the one exception: it needs the real, frozen eval set. Frozen means the same seeded rows every run, so a score change can only come from the model.
Two habits keep a suite trustworthy. First, know what not to test: never test sklearn's own maths, and never compare floats for exact equality — use a tolerance, or a threshold. Second, hunt flakiness: a flaky test — a test that passes and fails across runs with no code change — usually hides unseeded randomness. A team that reruns until green has stopped trusting red, and then the suite protects nobody. Seeds from config plus tolerances remove almost every flake in ML code.
The definition of done mirrors module 4.1's clone ritual. A teammate clones the repo, builds the venv, trains once, and runs pytest: every test passes. Then she plants any one of three bugs — a formula typo, a schema change, a worse model — and the right layer fails, naming the problem. A suite that passes on good code and fails on planted bugs is the alarm system module 4.6 will wire into CI.
Work inside churn-repo from 4.1 and 4.2, venv active, models/churn_pipeline.joblib trained. Install the two new tools and pin them the 4.1 way: pip install pytest httpx, then add both to requirements.txt (TestClient needs httpx).
Step 1 — the first unit test, worked out by hand. The tests/ folder has waited empty since 4.1. Its first citizen pins add_features to numbers you can verify on paper: income 24000, three orders at 600 gives spend_share (600 × 3 ÷ 3) ÷ 24000 = 0.025, and a last order on June 20 sits 10 days before the June 30 snapshot:
# tests/test_features.py — layer 2: the feature function, pinned by hand
import pandas as pd
import pytest
from src.features import add_features
def row(**overrides):
base = {"city": "Ludhiana", "monthly_income": 24000.0,
"orders_last_90d": 3, "avg_order_value": 600.0,
"last_order_date": pd.Timestamp("2026-06-20"), "churned": 0}
base.update(overrides)
return pd.DataFrame([base])
def test_known_customer_worked_by_hand():
out = add_features(row())
assert out.loc[0, "days_since_last_order"] == 10
assert out.loc[0, "spend_share"] == pytest.approx(0.025)
assert out.loc[0, "salary_week_order"] == 0
def test_zero_orders_means_zero_spend_share():
out = add_features(row(orders_last_90d=0))
assert out.loc[0, "spend_share"] == pytest.approx(0.0)
def test_input_frame_is_never_modified():
df = row()
before = df.copy()
add_features(df)
pd.testing.assert_frame_equal(df, before)
Three tests, three different promises. The first pins the happy path to hand-worked numbers — note pytest.approx, a tolerance instead of exact float equality. The second is the zero-orders edge case: a customer with no orders must score a spend_share of exactly zero, not an error. The third enforces the docstring's quiet promise from 4.1 — the input frame is never modified.
Step 2 — run it, then break it on purpose. pytest finds the tests by name; no registration, no configuration:
(.venv) $ pytest -q
... [100%]
3 passed in 0.41s
# now plant Vivek's bug: in src/features.py change / 3 to / 30
(.venv) $ pytest -q
F.. [100%]
=================== FAILURES ===================
_______ test_known_customer_worked_by_hand _______
def test_known_customer_worked_by_hand():
out = add_features(row())
assert out.loc[0, "days_since_last_order"] == 10
> assert out.loc[0, "spend_share"] == pytest.approx(0.025)
E assert 0.0025 == 0.025 ± 2.5e-08
tests/test_features.py:21: AssertionError
1 failed, 2 passed in 0.44s
Read the failure like a report, top to bottom. F.. is the scoreboard: first test failed, two passed. The > arrow marks the exact line that stopped the run. The E line is the heart: the left value 0.0025 is what your code produced, the right is what you promised, and the gap names the bug — output ten times too small, so hunt for a stray zero. The last line gives the file and line number. This failure took 0.44 seconds; the same bug took Vivek's team three weeks. Undo the edit and rerun before moving on: 3 passed.
Step 3 — layer 1, data validation, plus the shared fixtures. tests/conftest.py is a special file pytest reads automatically; fixtures defined there are available to every test file:
# tests/conftest.py — shared fixtures for the whole suite
import pandas as pd
import pytest
from src.data import load_customers
@pytest.fixture(scope="session")
def customers():
"""The real seeded world — loaded once for the whole run."""
return load_customers()
@pytest.fixture()
def tiny_customers():
"""Twelve hand-typed rows: six churned, six not, all five cities."""
rows = []
for i in range(12):
rows.append({
"city": ["Ludhiana", "Jalandhar", "Amritsar",
"Patiala", "Bathinda"][i % 5],
"monthly_income": 15000.0 + 2000 * i,
"orders_last_90d": (i * 3) % 9,
"avg_order_value": 300.0 + 50 * i,
"last_order_date": pd.Timestamp("2026-06-30")
- pd.Timedelta(days=7 * i),
"churned": i % 2,
})
return pd.DataFrame(rows)
# tests/test_data.py — layer 1: is the incoming data even legal?
import config
REQUIRED = ["city", "monthly_income", "orders_last_90d",
"avg_order_value", "last_order_date", "churned"]
def test_schema_required_columns_present(customers):
assert set(REQUIRED) <= set(customers.columns)
def test_ranges_are_physically_possible(customers):
assert (customers["monthly_income"] > 0).all()
assert customers["orders_last_90d"].between(0, 270).all()
assert (customers["avg_order_value"] > 0).all()
assert (customers["last_order_date"] <= config.SNAPSHOT_DATE).all()
def test_no_nulls_in_required_columns(customers):
assert customers[REQUIRED].notna().all().all()
def test_no_duplicate_rows(customers):
assert not customers.duplicated().any()
def test_cities_are_from_the_known_list(customers):
assert set(customers["city"]) <= set(config.CITIES)
The customers fixture loads the real seeded world once per run — any test that names it as a parameter receives it. The tiny_customers fixture is the hand-typed twelve-row world the pipeline tests will fit on: small enough to run in seconds, known enough to reason about. The five data tests are the door guards from 4.2's schema idea, now pointed at your own training data: columns, ranges, nulls, duplicates, legal cities.
Step 4 — the guards catch a planted bad row. Append one poisoned line to data/customers.csv — a Jalandhar customer with income −5000 — and rerun only layer 1:
# planted at the end of data/customers.csv:
# Jalandhar,-5000,2,450.0,2026-05-14,0 <- negative income
(.venv) $ pytest tests/test_data.py -q
.F... [100%]
=================== FAILURES ===================
________ test_ranges_are_physically_possible ________
> assert (customers["monthly_income"] > 0).all()
E assert False
1 failed, 4 passed in 0.58s
The schema test still passes — the columns are fine — but the range guard trips. assert False is terse, so the next move is one line in a Python shell: customers[customers["monthly_income"] <= 0] prints the guilty row. In a real pipeline this is the test that fires when an upstream export changes, a currency column arrives in paise, or a broken join manufactures duplicates. Delete the planted row and rerun: 5 passed.
Step 5 — layers 3 and 4, plus the cap. Three short files finish the suite. The contract tests fit on the tiny fixture; the quality gate loads the shipped joblib and scores it on the frozen eval set — the same seeded split train.py prints; the smoke test calls the 4.2 app through TestClient:
# tests/test_pipeline.py — layer 3: the contract, on tiny fixture data
import joblib
import numpy as np
from src.features import add_features
from src.pipeline import build_pipeline
def fit_tiny(tiny):
X = add_features(tiny).drop(columns="churned")
return build_pipeline().fit(X, tiny["churned"]), X
def test_predict_shape_and_range(tiny_customers):
pipe, X = fit_tiny(tiny_customers)
assert pipe.predict(X).shape == (12,)
assert set(pipe.predict(X)) <= {0, 1}
proba = pipe.predict_proba(X)[:, 1]
assert ((proba >= 0) & (proba <= 1)).all()
def test_save_load_round_trip_preserves_predictions(tiny_customers, tmp_path):
pipe, X = fit_tiny(tiny_customers)
joblib.dump(pipe, tmp_path / "pipe.joblib")
loaded = joblib.load(tmp_path / "pipe.joblib")
assert np.array_equal(loaded.predict(X), pipe.predict(X))
assert np.allclose(loaded.predict_proba(X), pipe.predict_proba(X))
# tests/test_quality.py — layer 4: the gate a worse model cannot pass
import joblib
from sklearn.metrics import recall_score
from sklearn.model_selection import train_test_split
import config
MIN_RECALL = 0.60 # agreed floor: below this, churners slip away
MIN_ACCURACY = 0.78
def frozen_eval_set(customers):
X = customers.drop(columns="churned")
y = customers["churned"]
_, X_test, _, y_test = train_test_split(
X, y, test_size=config.TEST_SIZE,
random_state=config.SEED, stratify=y)
return X_test, y_test
def test_recall_gate(customers):
X_test, y_test = frozen_eval_set(customers)
model = joblib.load(config.MODEL_PATH)
assert recall_score(y_test, model.predict(X_test)) >= MIN_RECALL
def test_accuracy_gate(customers):
X_test, y_test = frozen_eval_set(customers)
model = joblib.load(config.MODEL_PATH)
assert model.score(X_test, y_test) >= MIN_ACCURACY
# tests/test_api.py — the cap: does the whole service answer?
from fastapi.testclient import TestClient
from app.main import app
JALANDHAR = {"city": "Jalandhar", "monthly_income": 24000,
"orders_last_90d": 1, "avg_order_value": 650,
"last_order_date": "2026-03-02"}
def test_smoke_health_and_predict():
with TestClient(app) as client: # runs lifespan: model loads
assert client.get("/health").json() == {
"status": "ok", "model_loaded": True}
r = client.post("/predict", json=JALANDHAR)
assert r.status_code == 200
assert r.json()["prediction"] in (0, 1)
assert 0.0 <= r.json()["churn_prob"] <= 1.0
def test_smoke_bad_input_bounces_with_422():
with TestClient(app) as client:
bad = {**JALANDHAR, "monthly_income": -1}
assert client.post("/predict", json=bad).status_code == 422
Walk the choices. The contract tests never touch the real data: twelve rows prove the pipeline fits, predicts the right shape, keeps probabilities inside [0, 1], and survives a save and load round trip with identical predictions — the promise 4.2's lifespan loader depends on. The gate is different on purpose: it loads the shipped models/churn_pipeline.joblib — the exact artifact 4.2 serves — and holds it to floors on the frozen eval set. The seed makes that eval set identical on every machine, so the gate can never fail because the dice rolled differently. The smoke tests reuse everything 4.2 taught: TestClient starts the app in-process, the lifespan loads the model, and no uvicorn, port or second terminal is involved. Note what the smoke test asserts: the service answers correctly shaped — it says nothing about model quality. That is the gate's job.
Step 6 — the whole suite, then the gate earns its keep. First the happy run, then a deliberately worsened model:
(.venv) $ pytest -q
.............. [100%]
14 passed in 4.12s
# sabotage: in train.py, fit on the three RAW_FEATURES only,
# dropping the engineered signals — then retrain and overwrite:
(.venv) $ python train.py
accuracy: 0.752 recall: 0.487 caught: 19 of 39
saved: models/churn_pipeline.joblib
(.venv) $ pytest -q
............FF [100%]
E assert 0.4871... >= 0.6 # test_recall_gate
E assert 0.752 >= 0.78 # test_accuracy_gate
2 failed, 12 passed in 4.20s
This is the module's headline result — read it twice. The sabotaged train.py ran happily and saved a model that misses 20 of 39 churners; nothing in training complains, because training cannot know what "good enough" means. Twelve tests stay green — the data is legal, the features are right, the contract holds, the API answers politely — because a worse model breaks none of those promises. Only the quality gate goes red, and its E lines carry the evidence: recall 0.487 against a floor of 0.60. This is the test that fails when someone ships a worse model, whatever the reason — a lost feature, a bad merge, silently changed data. Restore train.py, retrain to 0.824, and the board is green again: 14 passed.
Work inside churn-repo, venv active, model trained. Every world is seeded: where a check states a number, your number must match it.
First test, first failure. Install pytest and httpx, pin both in requirements.txt, and build tests/test_features.py with the three tests from the watch section — work the 0.025 and the 10 days out on paper before you type them. Run pytest -q and record the output. Now plant the divisor bug (/ 3 to / 30 in src/features.py), rerun, and read the failure top to bottom. Restore the code and rerun. Check: the first run prints 3 passed; with the bug planted, exactly one test fails and the E line shows 0.0025 on the left and 0.025 on the right; you write three sentences naming what the F.. scoreboard, the > arrow line and the E line each told you; after restoring, 3 passed returns.
The full four-layer suite. Build the rest of the watch section's suite: tests/conftest.py with the customers and tiny_customers fixtures, tests/test_data.py (five guards: schema, ranges, nulls, duplicates, cities), tests/test_pipeline.py (shape and range, save/load round trip), tests/test_quality.py (recall floor 0.60, accuracy floor 0.78 on the frozen eval set), and tests/test_api.py (two smoke tests through TestClient). Commit the tests folder. Check: pytest -q prints 14 passed; each layer also passes alone (pytest tests/test_data.py and friends); deleting models/churn_pipeline.joblib makes the gate and smoke tests fail while layers 1-3 stay green, and python train.py heals the board; git ls-files shows five test files and no joblib.
The regression-catch drill. Prove the suite catches real damage by planting three bugs, one at a time, undoing each before the next. Bug A: the divisor typo in add_features (do not retrain). Bug B: a schema change — rename the orders_last_90d column to orders_90d in your data-building code, delete the csv so it regenerates. Bug C: the worse model — refit on raw features only and overwrite the joblib, the watch section's sabotage. After each plant, run pytest -q and record in regressions.md: the bug, which tests failed, which layer named the cause. Check: regressions.md shows Bug A caught by test_features with the two numbers in the E line; Bug B caught by the layer-1 schema guard naming the missing column (note in one sentence which other layers also fell, and why the data test is the one that names the cause); Bug C failing only the two gate tests while both smoke tests stay green — with one sentence on why a green smoke test must never be read as model quality; the drill ends with the board restored to 14 passed.
The flake hunt. Build tests/test_flaky_demo.py with one deliberately bad test: sample 100 rows from the customers fixture WITHOUT a random_state, fit the pipeline on them, and assert recall on the sample equals your exact full-data recall. Run pytest tests/test_flaky_demo.py ten times and record the pass/fail pattern. Now fix it the professional way: seed the sample with config.SEED and replace the exact-equality assert with a sensible floor. Run ten more times. Delete one more thing before finishing: write, then remove, a test that asserts sklearn's LogisticRegression stores a coefficient matrix — and record in one sentence why it never deserved to exist. Check: your log shows at least one failure in the ten unseeded runs and zero in the ten seeded runs; the final file uses a seed plus a threshold, no exact float equality; your sentence states that it tested sklearn's internals, not your promises; pytest -q on the full suite still ends green.
Three testing failures cover most of what this module meets at work. Each has a War Room entry with the full diagnosis path:
- WR-103 — the divisor that lied for three weeks. A feature formula typo shipped silently because nothing asserted the function's output; recall decayed until a monthly review caught it. Deterministic in-and-out unit tests with hand-worked numbers make that silence impossible.
- WR-104 — the suite nobody trusted. Unseeded randomness and exact float asserts made tests fail at random, the team learned to rerun until green, and then a real failure sailed past ignored. Seeds, tolerances and same-day quarantine of flakes restore the meaning of red.
- WR-105 — fourteen green tests, zero protection. Every test asserted only that the code ran without raising, so a much worse model passed the whole suite. Quality-gate thresholds on a frozen eval set — asserting on numbers, not survival — closed the hole.
ML code fails silently: a wrong number looks exactly like a right one, and no exception warns you. pytest turns your repo's promises into assertions that fail loudly, with the file, the line and both numbers in the report. The four layers split the job: data guards, feature tests with hand-worked numbers, pipeline-contract checks, and a model-quality gate with floors on a frozen eval set. A TestClient smoke test caps the stack by proving the API answers, while fixtures, seeds and tolerances keep the whole suite fast and honest. A suite earns trust when planted bugs — a formula typo, a schema change, a worse model — each turn the right layer red.
- Bank and NBFC model risk. Indian banks and NBFCs must show the regulator that a scoring model was checked before release, and their checks look like today's gate: a frozen eval set, agreed floors, and a recorded pass. UPI fraud models retrain on fresh data often, and the gate is what stops a worse retrain from reaching live payments.
- AI QA as a career. Indian GCCs and IT-services firms now post AI QA roles whose day is this module: pytest, data validation on ingest, and regression gates around models. The four-layer answer is what separates candidates in those interviews.
- Open-source ML libraries. scikit-learn and pandas run tens of thousands of seeded, tolerance-based tests on every proposed change. Today's habits are the same discipline at small scale.
Answer these before opening the gate assessment. Each question names its kind.
type=output-prediction— The 14-test suite is green. A teammate editssrc/features.pyso spend_share divides by 300 instead of 3, and does not retrain. Predict the nextpytest -qrun: which test fails first, what the E line shows for the hand-worked Ludhiana row, and what the two smoke tests report.type=debug— A fresh clone of your repo runspytest -qand gets 12 passed with the two gate tests failing onFileNotFoundError: models/churn_pipeline.joblib, while your laptop shows 14 passed on the same commit. Explain why the clone lacks a file your machine has, name the 4.1 decision that caused it, and give the one command that heals the clone.type=interpretation— A morning run prints.F...intests/test_data.py, failingtest_ranges_are_physically_possiblewithE assert Falseon the income line; yesterday it was green and no code changed. Read the situation: what does a data-layer red with unchanged code tell you, and what is your next command to find the evidence?type=scenario-judgement— Before release week a teammate proposes softeningtest_quality.pyto assert only thatmodel.predictruns without raising, because "the thresholds keep blocking us". Judge the proposal, name the War Room failure it recreates, and say what a blocked release is actually telling the team.type=trade-off— A recall floor of 0.66 (a whisker under the current 0.667) versus a floor of 0.50: give one concrete gain and one concrete cost of each, pick one for Vivek's NBFC service, and name the process you would attach to any gate failure so the floor stays meaningful.
Explain to a friend who cooks for a college mess why the kitchen checks food at four points — vegetables inspected at delivery, each masala mix tasted against the recipe, a trial plate checked before the full batch, and the final dish tasted against the standard before serving — instead of waiting for students to complain. Cover why a complaint three weeks later is the most expensive way to find a bad ingredient. Use at most 4 sentences and no technical word beyond "recipe". If your friend can say why each checkpoint catches a mistake the others would miss, you own this module.