Testing

Testing scikit-learn Models & Pipelines: Unit Tests, Cross-Validation, and CI Integration

HelpMeTest

HelpMeTest

— 6 min read
Testing scikit-learn Models & Pipelines: Unit Tests, Cross-Validation, and CI Integration

Most ML projects have zero tests for their scikit-learn pipelines. The model "works" if it trains without crashing and the accuracy metric looks reasonable. But this leaves a class of bugs invisible:

  • A custom transformer that silently drops rows
  • A Pipeline that applies preprocessing in the wrong order
  • A feature engineering function that leaks test set information
  • A model that fails on a new input shape at inference time

Testing scikit-learn code is straightforward once you know what to test. This guide covers the full spectrum: transformer unit tests, Pipeline contract tests, cross-validation testing, and CI setup.

What to Test in an ML Pipeline

Component What to test
Custom transformers fit, transform, output shape & dtype
Feature engineering functions Input/output contract, null handling
Preprocessing pipeline Correct order, no data leakage
Model training Trains without error, output shape
Prediction Input/output shape, value ranges
Serialization save → load → predict gives same result

Setup

pip install scikit-learn pytest pandas numpy

Testing Custom Transformers

Custom scikit-learn transformers must implement fit and transform. Test them with check_estimator plus your own tests:

# my_pipeline/transformers.py
import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin

class LogTransformer(BaseEstimator, TransformerMixin):
    """Apply log1p to specified numeric columns."""
    
    def __init__(self, columns=None):
        self.columns = columns
    
    def fit(self, X, y=None):
        if self.columns is None:
            self.columns_ = X.select_dtypes(include=np.number).columns.tolist()
        else:
            self.columns_ = self.columns
        return self
    
    def transform(self, X):
        X = X.copy()
        X[self.columns_] = np.log1p(X[self.columns_])
        return X
# tests/test_transformers.py
import pytest
import numpy as np
import pandas as pd
from sklearn.utils.estimator_checks import parametrize_with_checks
from my_pipeline.transformers import LogTransformer

@pytest.fixture
def sample_df():
    return pd.DataFrame({
        "revenue": [100.0, 200.0, 0.0, 50.0],
        "quantity": [2, 4, 0, 1],
        "product": ["A", "B", "A", "C"],  # non-numeric
    })

# Test the sklearn estimator interface compliance
@parametrize_with_checks([LogTransformer()])
def test_estimator_compliance(estimator, check):
    check(estimator)

def test_log_transformer_output_shape(sample_df):
    transformer = LogTransformer(columns=["revenue", "quantity"])
    result = transformer.fit_transform(sample_df)
    
    assert result.shape == sample_df.shape

def test_log_transformer_applies_log1p(sample_df):
    transformer = LogTransformer(columns=["revenue"])
    result = transformer.fit_transform(sample_df)
    
    expected = np.log1p([100.0, 200.0, 0.0, 50.0])
    np.testing.assert_array_almost_equal(result["revenue"].values, expected)

def test_log_transformer_preserves_non_numeric(sample_df):
    transformer = LogTransformer(columns=["revenue"])
    result = transformer.fit_transform(sample_df)
    
    assert list(result["product"]) == list(sample_df["product"])

def test_log_transformer_handles_zeros(sample_df):
    """log1p(0) = 0, not -inf. Should not raise."""
    transformer = LogTransformer(columns=["quantity"])
    result = transformer.fit_transform(sample_df)
    
    assert result["quantity"].iloc[2] == 0.0  # log1p(0) = 0

def test_log_transformer_fit_then_transform(sample_df):
    """fit and transform called separately should give same result as fit_transform."""
    transformer = LogTransformer(columns=["revenue"])
    
    transformer.fit(sample_df)
    result_separate = transformer.transform(sample_df)
    result_combined = LogTransformer(columns=["revenue"]).fit_transform(sample_df)
    
    pd.testing.assert_frame_equal(result_separate, result_combined)

def test_log_transformer_auto_detects_numeric_columns(sample_df):
    """When columns=None, should auto-select numeric columns."""
    transformer = LogTransformer()
    result = transformer.fit_transform(sample_df)
    
    # "product" column should be unchanged
    assert list(result["product"]) == list(sample_df["product"])
    # Numeric columns should be transformed
    assert result["revenue"].iloc[0] != sample_df["revenue"].iloc[0]

Testing sklearn Pipelines

Pipeline objects chain multiple steps. Test the full pipeline's input/output contract:

# my_pipeline/pipeline.py
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from my_pipeline.transformers import LogTransformer

def build_churn_pipeline():
    return Pipeline([
        ("log_transform", LogTransformer(columns=["revenue", "tenure_days"])),
        ("scaler", StandardScaler()),
        ("model", LogisticRegression(random_state=42)),
    ])
# tests/test_pipeline.py
import pytest
import numpy as np
import pandas as pd
from sklearn.datasets import make_classification
from my_pipeline.pipeline import build_churn_pipeline

@pytest.fixture
def training_data():
    """Minimal labeled dataset for pipeline tests."""
    X = pd.DataFrame({
        "revenue": [100, 200, 50, 300, 150, 75, 400, 80],
        "tenure_days": [30, 90, 10, 180, 60, 5, 365, 45],
        "support_tickets": [1, 0, 3, 0, 1, 5, 0, 2],
    })
    y = np.array([0, 0, 1, 0, 0, 1, 0, 1])
    return X, y

def test_pipeline_fits_without_error(training_data):
    X, y = training_data
    pipeline = build_churn_pipeline()
    pipeline.fit(X, y)  # Should not raise

def test_pipeline_predict_returns_correct_shape(training_data):
    X, y = training_data
    pipeline = build_churn_pipeline()
    pipeline.fit(X, y)
    
    predictions = pipeline.predict(X)
    assert predictions.shape == (len(X),)

def test_pipeline_predict_returns_binary_labels(training_data):
    X, y = training_data
    pipeline = build_churn_pipeline()
    pipeline.fit(X, y)
    
    predictions = pipeline.predict(X)
    assert set(predictions).issubset({0, 1})

def test_pipeline_predict_proba_returns_probabilities(training_data):
    X, y = training_data
    pipeline = build_churn_pipeline()
    pipeline.fit(X, y)
    
    probas = pipeline.predict_proba(X)
    assert probas.shape == (len(X), 2)
    assert (probas >= 0).all()
    assert (probas <= 1).all()
    np.testing.assert_array_almost_equal(probas.sum(axis=1), np.ones(len(X)))

def test_pipeline_single_row_prediction(training_data):
    """Critical: inference should work on a single row."""
    X, y = training_data
    pipeline = build_churn_pipeline()
    pipeline.fit(X, y)
    
    single_row = X.iloc[[0]]  # Preserve DataFrame shape (not Series)
    prediction = pipeline.predict(single_row)
    assert prediction.shape == (1,)

def test_pipeline_step_order():
    """Verify steps are in the expected order."""
    pipeline = build_churn_pipeline()
    step_names = [name for name, _ in pipeline.steps]
    
    assert step_names == ["log_transform", "scaler", "model"]

Testing for Data Leakage

Data leakage is when test set information influences training. A common source: fitting a scaler on the full dataset, then splitting. Test for this:

def test_no_data_leakage_in_pipeline(training_data):
    """Pipeline should fit only on training data, not test data."""
    from sklearn.model_selection import train_test_split
    
    X, y = training_data
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=42)
    
    pipeline = build_churn_pipeline()
    pipeline.fit(X_train, y_train)
    
    # Scaler fitted stats should come from training set only
    scaler = pipeline.named_steps["scaler"]
    train_mean = X_train["revenue"].mean()
    
    # Pipeline's internal scaler mean should approximate training set mean
    # (after log transform, but the principle holds)
    assert scaler.mean_ is not None  # Was fitted
    
    # Key test: predict on test set should not raise
    predictions = pipeline.predict(X_test)
    assert predictions.shape == (len(X_test),)

Testing Cross-Validation

Don't test a single CV run — test that your CV harness is set up correctly:

from sklearn.model_selection import cross_val_score, StratifiedKFold

def test_cross_validation_runs_without_error(training_data):
    X, y = training_data
    pipeline = build_churn_pipeline()
    
    cv = StratifiedKFold(n_splits=3, shuffle=True, random_state=42)
    scores = cross_val_score(pipeline, X, y, cv=cv, scoring="roc_auc")
    
    assert scores.shape == (3,)
    assert (scores >= 0).all()
    assert (scores <= 1).all()

def test_model_performance_above_baseline(training_data):
    """Model should beat a naive majority-class baseline."""
    X, y = training_data
    pipeline = build_churn_pipeline()
    
    cv = StratifiedKFold(n_splits=3, shuffle=True, random_state=42)
    scores = cross_val_score(pipeline, X, y, cv=cv, scoring="accuracy")
    
    majority_class_baseline = max(y.mean(), 1 - y.mean())
    assert scores.mean() >= majority_class_baseline - 0.05  # Allow 5% slack
Note: Don't assert exact accuracy values in unit tests — they depend on random seeds and data ordering. Test that the CV harness runs and that scores are valid numbers in [0, 1].

Testing Model Serialization

Models must survive save/load cycles for production deployment:

import pickle
import tempfile
from pathlib import Path

def test_pipeline_serialization(training_data):
    X, y = training_data
    pipeline = build_churn_pipeline()
    pipeline.fit(X, y)
    
    with tempfile.TemporaryDirectory() as tmpdir:
        model_path = Path(tmpdir) / "model.pkl"
        
        # Save
        with open(model_path, "wb") as f:
            pickle.dump(pipeline, f)
        
        # Load
        with open(model_path, "rb") as f:
            loaded_pipeline = pickle.load(f)
        
        # Predictions should be identical
        original_predictions = pipeline.predict(X)
        loaded_predictions = loaded_pipeline.predict(X)
        
        np.testing.assert_array_equal(original_predictions, loaded_predictions)

For production, use joblib (faster for large arrays):

import joblib

joblib.dump(pipeline, "model.joblib")
loaded = joblib.load("model.joblib")

Testing Feature Importance (for Interpretable Models)

def test_feature_importances_available(training_data):
    """For tree-based models, feature importances should be accessible after fit."""
    from sklearn.ensemble import RandomForestClassifier
    from my_pipeline.transformers import LogTransformer
    
    pipeline = Pipeline([
        ("log_transform", LogTransformer(columns=["revenue", "tenure_days"])),
        ("model", RandomForestClassifier(n_estimators=10, random_state=42)),
    ])
    
    X, y = training_data
    pipeline.fit(X, y)
    
    importances = pipeline.named_steps["model"].feature_importances_
    assert importances.shape == (X.shape[1],)
    assert abs(importances.sum() - 1.0) < 1e-6  # Should sum to 1

CI Configuration

# .github/workflows/ml-tests.yml
name: ML Pipeline Tests

on: [push, pull_request]

jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with: { python-version: "3.12" }
      - run: pip install scikit-learn pandas numpy pytest pytest-cov
      - run: pytest tests/ -v --cov=my_pipeline --cov-report=xml
      - uses: codecov/codecov-action@v4

Regression Testing for Model Accuracy

When you retrain, compare against a saved baseline to detect unexpected performance degradation:

# tests/test_regression.py
import json
from pathlib import Path

BASELINE_PATH = Path("tests/fixtures/performance_baseline.json")

def test_model_performance_matches_baseline(training_data):
    from sklearn.model_selection import cross_val_score, StratifiedKFold
    
    X, y = training_data
    pipeline = build_churn_pipeline()
    
    cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    score = cross_val_score(pipeline, X, y, cv=cv, scoring="roc_auc").mean()
    
    if BASELINE_PATH.exists():
        with open(BASELINE_PATH) as f:
            baseline = json.load(f)
        
        # Allow up to 2% degradation
        assert score >= baseline["roc_auc"] - 0.02, (
            f"Performance degraded: {score:.4f} vs baseline {baseline['roc_auc']:.4f}"
        )
    else:
        # First run: save baseline
        BASELINE_PATH.parent.mkdir(exist_ok=True)
        with open(BASELINE_PATH, "w") as f:
            json.dump({"roc_auc": float(score)}, f)
        print(f"Baseline saved: AUC = {score:.4f}")

What Not to Test

  • Don't test sklearn internals — StandardScaler works. You're testing your code.
  • Don't hardcode exact probabilities — they change with random seeds.
  • Don't test that accuracy is above 90% on toy data — that's not meaningful for production performance.

Focus on contracts: shape, dtype, valid ranges, serialization round-trips, and the absence of silent failures.

Conclusion

Testing scikit-learn code is about enforcing contracts, not validating that the algorithm does what the algorithm always does. Test your custom transformers with parametrize_with_checks, test Pipeline step order and input/output shapes, test serialization round-trips, and add lightweight regression tests for performance baselines. These tests catch the class of bugs that actually occur in production ML pipelines — and they run in seconds.

Read more

Start now free