Testing ML Systems
Tests for data, code, and models.
Software engineers test their code before shipping. Machine learning needs the same discipline, and more. You must test the code, the data, and the model itself. A good test suite is the guard at the gate of your registry. In this lesson, we write real tests with pytest.

- Test the code, the data, the model
- Automatic, on every change
- A guard at the registry gate
Three kinds of ML tests
There are three kinds of tests. Data tests check the inputs, no missing values, sensible ranges, a plausible share of churners. Model tests check quality, is the model better than a trivial baseline, and better than the current champion? And behaviour tests check common sense. If a customer calls support more often, their churn risk should not go down. These tests catch problems that a single accuracy number hides.
Model tests
Here are model tests. The first compares our model with a baseline, that always predicts the most common answer. We use the F one score, because only about one customer in eight leaves, so a lazy model that says nobody leaves, already scores high accuracy. The model must beat the baseline clearly. The second test is a behaviour test. Two identical customers, except that one called support six times. The riskier one must receive a higher churn probability. And the third checks that outputs are valid probabilities.
import pandas as pd
from sklearn.dummy import DummyClassifier
from sklearn.metrics import f1_score
from churn import make_data
def test_beats_the_baseline(model):
test = make_data(1000, seed=2)
X, y = test.drop(columns="churn"), test["churn"]
baseline = DummyClassifier(strategy="most_frequent").fit(X, y)
assert f1_score(y, model.predict(X)) > f1_score(y, baseline.predict(X)) + 0.2
def test_more_support_calls_never_lowers_risk(model):
low = pd.DataFrame([{"tenure": 24, "charges": 70.0, "calls": 0, "contract": 0}])
high = pd.DataFrame([{"tenure": 24, "charges": 70.0, "calls": 6, "contract": 0}])
assert model.predict_proba(high)[0][1] > model.predict_proba(low)[0][1]
def test_prediction_is_a_probability(model):
p = model.predict_proba(pd.DataFrame([{"tenure": 12, "charges": 50.0, "calls": 2, "contract": 1}]))[0][1]
assert 0.0 <= p <= 1.0Running the tests
Let's run the suite. Pytest finds the tests, runs them, and prints a dot for every pass. All six pass. Now let's inject bad data on purpose, a negative monthly charge. We run the data tests again. This time, a test fails, and pytest shows exactly which check broke, and the offending value. In a real pipeline, this failure would stop the process, before a bad model ever reached the registry.
pytest -q -p no:cacheprovider -W ignore ...... [100%] 6 passed in 0.05s INJECT_BAD_DATA=1 pytest -q -p no:cacheprovider -W ignore tests/test_data.py 2>&1 | grep -E 'assert|passed|failed' | head -4 > assert (data["charges"] > 0).all() E assert np.False_ FAILED tests/test_data.py::test_charges_are_positive - assert np.False_ 1 failed, 2 passed in 0.08s
File used: conftest.py
import os
import pytest
from sklearn.linear_model import LogisticRegression
from churn import make_data
@pytest.fixture(scope="session")
def data():
df = make_data(3000, seed=1)
if os.environ.get("INJECT_BAD_DATA"):
df.loc[0, "charges"] = -5.0
return df
@pytest.fixture(scope="session")
def model(data):
X, y = data.drop(columns="churn"), data["churn"]
return LogisticRegression(max_iter=500).fit(X, y)File used: churn.py
import numpy as np, pandas as pd
def make_data(n, seed, shifted=False):
rng = np.random.default_rng(seed)
tenure = rng.integers(1, 72, n)
charges = rng.normal(65 + (28 if shifted else 0), 18, n).clip(20, 140).round(2)
calls = rng.poisson(1.5 + (1.5 if shifted else 0), n)
contract = rng.choice([0, 1, 2], n, p=[.55, .25, .20] if not shifted else [.30, .30, .40])
logit = -1.0 - 0.04 * tenure + 0.025 * (charges - 65) + 0.45 * calls - 0.8 * contract
if shifted:
logit = -0.2 - 0.01 * tenure - 0.01 * (charges - 65) + 0.2 * calls - 0.2 * contract
churn = (rng.random(n) < 1 / (1 + np.exp(-logit))).astype(int)
return pd.DataFrame({"tenure": tenure, "charges": charges, "calls": calls, "contract": contract, "churn": churn})File used: tests/test_data.py
def test_no_missing_values(data):
assert not data.isna().any().any()
def test_charges_are_positive(data):
assert (data["charges"] > 0).all()
def test_churn_rate_is_plausible(data):
assert 0.1 < data["churn"].mean() < 0.6File used: tests/test_model.py
import pandas as pd
from sklearn.dummy import DummyClassifier
from sklearn.metrics import f1_score
from churn import make_data
def test_beats_the_baseline(model):
test = make_data(1000, seed=2)
X, y = test.drop(columns="churn"), test["churn"]
baseline = DummyClassifier(strategy="most_frequent").fit(X, y)
assert f1_score(y, model.predict(X)) > f1_score(y, baseline.predict(X)) + 0.2
def test_more_support_calls_never_lowers_risk(model):
low = pd.DataFrame([{"tenure": 24, "charges": 70.0, "calls": 0, "contract": 0}])
high = pd.DataFrame([{"tenure": 24, "charges": 70.0, "calls": 6, "contract": 0}])
assert model.predict_proba(high)[0][1] > model.predict_proba(low)[0][1]
def test_prediction_is_a_probability(model):
p = model.predict_proba(pd.DataFrame([{"tenure": 12, "charges": 50.0, "calls": 2, "contract": 1}]))[0][1]
assert 0.0 <= p <= 1.0Accuracy can lie
One more lesson from that baseline test. Accuracy can lie. When most customers stay, a model that always says stay, is right most of the time, and totally useless. For imbalanced problems, use precision, recall, and F one, or the area under the curve. Always compare against a trivial baseline, and always ask who is being missed. A test that includes a baseline, is a test that cannot be fooled by a flattering number.
Accuracy alone
- Hides class imbalance
- A lazy model scores high
- Looks great, helps nobody
Better metrics
- Precision, recall, F1
- Compare against a baseline
- Look at who is missed
Summary
Excellent. Test the data, the quality of the model, and its behaviour. Beware of accuracy on imbalanced problems, and always compare with a baseline. These tests guard the registry gate. Next, getting the model out of Python, into portable formats.
Key takeaways- Data, quality, behaviour tests
- Compare with a baseline
- Use F1 on imbalanced data