Testing ML Code & Data
Writing unit tests, integration tests, behavioral tests, and data validation suites for reliable machine learning code.
The Silent Failure Problem
In standard software development, bugs cause explicit exceptions: NullPointerException or SyntaxError.
In machine learning development, code can run cleanly without any error while producing garbage predictions:
- Normalizing pixels by $255.0$ twice.
- Feature column order getting silently swapped.
- Missing values filled with zero instead of column median.
To catch silent failures, machine learning testing requires a Three Tier Testing Strategy.
┌─────────────────────────────────────────────────────────────┐
│ 1. CODE UNIT TESTS: Feature functions, tensor shapes. │
│ 2. DATA QUALITY TESTS: Schema types, nulls, range bounds. │
│ 3. MODEL BEHAVIOR TESTS: Directionality, edge cases, noise. │
└─────────────────────────────────────────────────────────────┘
1. Code Unit Tests (pytest)
Test deterministic pipeline functions in isolation using small synthetic mock datasets:
- Shape Tests: Verify feature transformation functions output expected tensor shapes (
[Batch_Size, Feature_Dim]). - Preprocessing Invariants: Confirm text normalization functions lowercase strings and trim whitespace consistently.
- Nan Guards: Assert custom loss functions never return
NaNorInfinityvalues.
2. Data Quality Tests (Great Expectations)
Validate incoming training and serving datasets before executing expensive model operations:
- Schema Constraints: Confirm required column names and datatypes exist.
- Value Bounds: Assert user age values fall between $0$ and $120$.
- Null Thresholds: Fail data ingestion if missing values in critical columns exceed $1$ percent.
3. Model Behavioral Testing (CheckList Framework)
Evaluate trained model behaviors using three functional test patterns:
┌──────────────────────────┬──────────────────────────┬──────────────────────────┐
│ A. INVARIANCE TESTS │ B. DIRECTIONAL TESTS │ C. MINIMUM FUNCTIONALITY │
├──────────────────────────┼──────────────────────────┼──────────────────────────┤
│ Changing non semantic │ Increasing positive │ Model correctly handles │
│ input details (changing │ input signals (adding │ obvious benchmark edge │
│ user name from John to │ positive reviews) MUST │ cases (negations like │
│ Mary) MUST NOT change │ INCREASE predicted sentiment| "not good"). │
│ model prediction score. │ score. │ │
└──────────────────────────┴──────────────────────────┴──────────────────────────┘
Say this out loud
Testing machine learning code requires testing software functions, data quality, and model behavioral predictions. Unit tests check feature engineering logic and tensor shapes. Data tests validate schemas and value bounds using tools like Great Expectations. Behavioral tests evaluate invariance, directional expectations, and minimum functionality to prevent silent production failures.
Followups to expect
- What is Integration Testing in ML? Testing the full end to end pipeline on a small synthetic dataset to ensure data loading, feature extraction, model prediction, and post processing execute cleanly together.
- What is Shadow Testing? Deploying a candidate model alongside the live production model to evaluate real world predictions on production traffic without displaying results to end users.
Check yourself
Why do standard software unit tests fail to catch many critical errors in machine learning applications?