Testing Scikit-learn Applications: Unit Tests to Integration

Testing Scikit-learn Applications: Unit Tests to Integration

Machine learning applications are notoriously difficult to test. Unlike traditional software, where outputs are deterministic and easily predictable, ML systems involve probabilistic models, data-dependent behavior, and subtle numerical edge cases. Scikit-learn, the most popular Python library for classical machine learning, provides a rich ecosystem not only for building models but also for testing them. This tutorial walks you through a complete testing strategy — from unit-testing individual transformers to integration-testing entire ML pipelines.

Why Testing ML Code Matters

Many data scientists skip rigorous testing because notebooks feel exploratory. But once a model moves toward production, untested code becomes a liability. Bugs in feature engineering can silently corrupt predictions. A refactored pipeline can change model behavior without anyone noticing. Tests give you confidence that:

• Custom transformers behave as expected on edge cases.
• Pipelines remain reproducible across library versions.
• Model performance does not regress after refactoring.
• Data preprocessing handles missing values and unseen categories gracefully.
• Retrained models meet minimum quality thresholds before deployment.

What to Test in a Scikit-learn Application

A typical scikit-learn application has several layers, each requiring a different testing approach:

• Unit tests — individual components like custom transformers, metric functions, and data loaders.
• Component tests — estimators in isolation, verifying fit/predict contracts.
• Integration tests — full pipelines end-to-end, including preprocessing, training, and inference.
• Regression tests — snapshots of model behavior to catch unintended changes.
• Performance tests — guards ensuring metrics stay above acceptable thresholds.

Setting Up the Project

Before writing tests, set up a clean project structure. We will use pytest as the test runner because of its concise syntax and powerful fixtures.

ml_project/
├── src/
│   └── ml_project/
│       ├── __init__.py
│       ├── transformers.py
│       ├── pipeline.py
│       └── train.py
├── tests/
│   ├── __init__.py
│   ├── conftest.py
│   ├── test_transformers.py
│   ├── test_pipeline.py
│   └── test_integration.py
├── requirements.txt
└── pyproject.toml

Install the required dependencies:

pip install scikit-learn pytest pandas numpy

Building a Custom Transformer

To have something meaningful to test, let's build a custom scikit-learn transformer that adds an engineered feature. This is a common real-world scenario where bugs creep in.

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


class RatioFeatureAdder(BaseEstimator, TransformerMixin):
    """Adds a ratio between two numeric columns."""

    def __init__(self, numerator: str, denominator: str, new_column: str):
        self.numerator = numerator
        self.denominator = denominator
        self.new_column = new_column

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        X = X.copy()
        denom = X[self.denominator].replace(0, np.nan)
        X[self.new_column] = X[self.numerator] / denom
        return X

Writing Unit Tests for Transformers

Unit tests verify that the transformer behaves correctly in isolation. We test the happy path, division by zero, and immutability of input data.

# tests/test_transformers.py
import numpy as np
import pandas as pd
import pytest

from ml_project.transformers import RatioFeatureAdder


@pytest.fixture
def sample_df():
    return pd.DataFrame({
        "revenue": [100, 200, 300],
        "cost": [10, 20, 0],
    })


def test_ratio_feature_added_correctly(sample_df):
    adder = RatioFeatureAdder("revenue", "cost", "margin")
    result = adder.transform(sample_df)

    assert "margin" in result.columns
    assert result.loc[0, "margin"] == 10.0
    assert result.loc[1, "margin"] == 10.0


def test_division_by_zero_produces_nan(sample_df):
    adder = RatioFeatureAdder("revenue", "cost", "margin")
    result = adder.transform(sample_df)

    assert np.isnan(result.loc[2, "margin"])


def test_transform_does_not_mutate_input(sample_df):
    adder = RatioFeatureAdder("revenue", "cost", "margin")
    original = sample_df.copy()
    adder.transform(sample_df)

    pd.testing.assert_frame_equal(sample_df, original)


def test_fit_returns_self(sample_df):
    adder = RatioFeatureAdder("revenue", "cost", "margin")
    assert adder.fit(sample_df) is adder

These tests cover the core contract of a scikit-learn transformer: fit returns self, transform returns the modified data, and inputs are not mutated.

Using Scikit-learn's Built-in Estimator Checks

Scikit-learn ships with a powerful utility called check_estimator that runs a standardized suite of tests against any estimator. This verifies API compliance, clone behavior, parameter handling, and many edge cases automatically.

# tests/test_transformers.py (continued)
from sklearn.utils.estimator_checks import check_estimator

from ml_project.transformers import RatioFeatureAdder


def test_estimator_compliance():
    # check_estimator runs dozens of API contract tests
    check_estimator(RatioFeatureAdder(
        numerator="revenue", denominator="cost", new_column="margin"
    ))

If your transformer follows the scikit-learn API correctly, this single test replaces dozens of hand-written checks. It validates that get_params, set_params, clone, and pipeline compatibility all work as expected.

Testing a Full Pipeline

Now let's assemble a pipeline that combines preprocessing, our custom transformer, and a model. Integration tests verify that all pieces work together.

# src/ml_project/pipeline.py
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

from ml_project.transformers import RatioFeatureAdder


def build_pipeline(numeric_features, categorical_features):
    numeric_preprocessing = Pipeline([
        ("imputer", SimpleImputer(strategy="median")),
        ("scaler", StandardScaler()),
    ])

    preprocessor = ColumnTransformer([
        ("num", numeric_preprocessing, numeric_features),
        ("cat", "passthrough", categorical_features),
    ])

    return Pipeline([
        ("ratio_adder", RatioFeatureAdder(
            numerator="revenue", denominator="cost", new_column="margin"
        )),
        ("preprocessor", preprocessor),
        ("classifier", LogisticRegression(max_iter=1000)),
    ])

Writing Integration Tests

Integration tests exercise the pipeline end-to-end. They use small synthetic datasets so they run fast and remain deterministic.

# tests/test_integration.py
import numpy as np
import pandas as pd
import pytest
from sklearn.datasets import make_classification
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split

from ml_project.pipeline import build_pipeline


@pytest.fixture
def synthetic_data():
    X, y = make_classification(
        n_samples=500, n_features=5, n_informative=3,
        random_state=42
    )
    df = pd.DataFrame(X, columns=["revenue", "cost", "f3", "f4", "f5"])
    df["revenue"] = df["revenue"] * 100 + 200
    df["cost"] = df["cost"] * 50 + 100
    return df, y


def test_pipeline_fits_and_predicts(synthetic_data):
    X, y = synthetic_data
    pipe = build_pipeline(
        numeric_features=["revenue", "cost", "f3", "f4", "f5", "margin"],
        categorical_features=[],
    )

    pipe.fit(X, y)
    predictions = pipe.predict(X)

    assert predictions.shape == (500,)
    assert set(np.unique(predictions)).issubset({0, 1})


def test_pipeline_meets_minimum_accuracy(synthetic_data):
    X, y = synthetic_data
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )

    pipe = build_pipeline(
        numeric_features=["revenue", "cost", "f3", "f4", "f5", "margin"],
        categorical_features=[],
    )
    pipe.fit(X_train, y_train)
    acc = accuracy_score(y_test, pipe.predict(X_test))

    assert acc > 0.7, f"Accuracy {acc:.3f} below threshold"


def test_pipeline_handles_missing_values():
    df = pd.DataFrame({
        "revenue": [100, 200, np.nan, 400],
        "cost": [10, 20, 30, 40],
        "f3": [1, 2, 3, 4],
        "f4": [5, 6, 7, 8],
        "f5": [9, 10, 11, 12],
    })
    y = np.array([0, 1, 0, 1])

    pipe = build_pipeline(
        numeric_features=["revenue", "cost", "f3", "f4", "f5", "margin"],
        categorical_features=[],
    )
    pipe.fit(df, y)
    preds = pipe.predict(df)

    assert preds.shape == (4,)

Shared Fixtures with conftest.py

To avoid duplicating fixture definitions across test files, move shared fixtures into conftest.py. Pytest automatically discovers them.

# tests/conftest.py
import numpy as np
import pandas as pd
import pytest
from sklearn.datasets import make_classification


@pytest.fixture
def synthetic_data():
    X, y = make_classification(
        n_samples=500, n_features=5, n_informative=3, random_state=42
    )
    df = pd.DataFrame(X, columns=["revenue", "cost", "f3", "f4", "f5"])
    df["revenue"] = df["revenue"] * 100 + 200
    df["cost"] = df["cost"] * 50 + 100
    return df, y


@pytest.fixture
def small_classification_data():
    return make_classification(
        n_samples=100, n_features=4, random_state=0
    )

Regression Testing with Model Snapshots

Regression tests ensure that refactoring does not silently change predictions. The idea is to save a snapshot of predictions from a known-good model and compare future runs against it.

# tests/test_pipeline.py
import numpy as np
import pytest

from ml_project.pipeline import build_pipeline


EXPECTED_PREDICTIONS = np.array([
    0, 1, 0, 1, 0, 1, 0, 1, 0, 1
])


def test_predictions_are_stable(synthetic_data):
    X, y = synthetic_data
    pipe = build_pipeline(
        numeric_features=["revenue", "cost", "f3", "f4", "f5", "margin"],
        categorical_features=[],
    )
    pipe.fit(X, y)
    preds = pipe.predict(X.head(10))

    np.testing.assert_array_equal(preds, EXPECTED_PREDICTIONS)

When you intentionally change the model, update the snapshot. When you do not, the test catches unintended drift.

Testing Hyperparameter Tuning

Grid search and randomized search introduce their own failure modes. Test that the search completes and returns a model that performs at least as well as the default.

# tests/test_integration.py (continued)
from sklearn.model_selection import GridSearchCV


def test_grid_search_improves_or_matches(synthetic_data):
    X, y = synthetic_data
    pipe = build_pipeline(
        numeric_features=["revenue", "cost", "f3", "f4", "f5", "margin"],
        categorical_features=[],
    )

    param_grid = {
        "classifier__C": [0.01, 0.1, 1.0, 10.0],
    }

    search = GridSearchCV(pipe, param_grid, cv=3, scoring="accuracy")
    search.fit(X, y)

    default_score = pipe.fit(X, y).score(X, y)
    best_score = search.best_score_

    assert best_score >= default_score - 0.05
    assert "classifier__C" in search.best_params_

Best Practices for Testing Scikit-learn Applications

• Use fixed random seeds. Always pass random_state to estimators, data splitters, and data generators so tests are reproducible.
• Keep test datasets small. Synthetic data with a few hundred rows runs in milliseconds and avoids coupling tests to external data files.
• Test edge cases explicitly. Empty DataFrames, all-missing columns, single-row inputs, and unseen categories during inference are common sources of production failures.
• Leverage check_estimator. For custom estimators, this built-in suite catches API violations you would never think to test manually.
• Separate unit from integration tests. Tag integration tests with @pytest.mark.integration so you can run fast unit tests locally and the full suite in CI.
• Assert on shapes and types, not just values. A model that returns the right accuracy but the wrong shape will break downstream consumers.
• Pin dependency versions in CI. Scikit-learn minor releases can change default behaviors; lock versions and test upgrades deliberately.
• Snapshots over thresholds where possible. Exact prediction snapshots are more sensitive than accuracy thresholds, but update them consciously.
• Test serialization. Verify that joblib.dump and joblib.load round-trip your pipeline without error.

Testing Model Serialization

Production models are almost always serialized. A common bug is a pipeline that trains fine but fails to load in a different environment. Add a round-trip test.

# tests/test_integration.py (continued)
import joblib
import os
import tempfile


def test_pipeline_serializes_and_loads(synthetic_data):
    X, y = synthetic_data
    pipe = build_pipeline(
        numeric_features=["revenue", "cost", "f3", "f4", "f5", "margin"],
        categorical_features=[],
    )
    pipe.fit(X, y)
    original_preds = pipe.predict(X)

    with tempfile.TemporaryDirectory() as tmp:
        path = os.path.join(tmp, "model.joblib")
        joblib.dump(pipe, path)
        loaded = joblib.load(path)

    loaded_preds = loaded.predict(X)
    np.testing.assert_array_equal(original_preds, loaded_preds)

Running the Test Suite

Run all tests with a single command:

pytest tests/ -v

To run only fast unit tests and skip integration tests:

pytest tests/ -v -m "not integration"

Add markers in pyproject.toml or pytest.ini:

[tool.pytest.ini_options]
markers = [
    "integration: marks tests as integration tests (deselect with '-m \"not integration\"')",
]

Conclusion

Testing scikit-learn applications requires a layered strategy that respects the unique challenges of machine learning code. Unit tests validate individual transformers and estimators, scikit-learn's check_estimator enforces API compliance, integration tests confirm that pipelines work end-to-end, and regression tests guard against silent behavioral drift. By combining synthetic fixtures, deterministic seeds, snapshot comparisons, and serialization round-trips, you build a safety net that lets you refactor, retrain, and ship models with confidence. The upfront investment in tests pays off the moment a refactor would have introduced a subtle bug — and the test suite catches it before your users do.