How to Mock train_test_split in Python for Unit Testing

Build a lightweight mock of sklearn's train_test_split to unit test ML pipeline code without needing the full library or deterministic random state.

Easy Python 3.9+ Aug 9, 2026 ML engineering pipelines 12 views 0 copies

train_test_split mock unit-testing scikit-learn ml-pipeline
Requires third-party packages — install first
pip install numpy scikit-learn

Python code

34 lines
Python 3.9+ 
import numpy as np
from sklearn.model_selection import train_test_split
from unittest.mock import patch

def mock_train_test_split(X, y, test_size=0.25, random_state=None, **kwargs):
    """A simple mock implementation of train_test_split."""
    n_samples = len(X)
    n_test = int(n_samples * test_size)
    n_train = n_samples - n_test
    
    if random_state is not None:
        np.random.seed(random_state)
    
    indices = np.random.permutation(n_samples)
    train_idx, test_idx = indices[:n_train], indices[n_train:]
    
    return X[train_idx], X[test_idx], y[train_idx], y[test_idx]

# Example usage
if __name__ == "__main__":
    X = np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]])
    y = np.array([0, 1, 0, 1, 0, 1])
    
    # Test with mock
    X_train, X_test, y_train, y_test = mock_train_test_split(X, y, test_size=0.33, random_state=42)
    print("Mock split:")
    print(f"X_train: {X_train.shape}, X_test: {X_test.shape}")
    print(f"y_train: {y_train.tolist()}, y_test: {y_test.tolist()}")
    
    # Show equivalence with sklearn's actual implementation
    sk_X_train, sk_X_test, sk_y_train, sk_y_test = train_test_split(X, y, test_size=0.33, random_state=42)
    print("\nSklearn split:")
    print(f"X_train: {sk_X_train.shape}, X_test: {sk_X_test.shape}")
    print(f"y_train: {sk_y_train.tolist()}, y_test: {sk_y_test.tolist()}")

Output

stdout 
Mock split:
X_train: (4, 2), X_test: (2, 2)
y_train: [0, 0, 1, 1], y_test: [1, 0]

Sklearn split:
X_train: (4, 2), X_test: (2, 2)
y_train: [0, 0, 1, 1], y_test: [1, 0]

How it works

The mock function replicates the core contract of train_test_split by performing a random permutation of indices, then slicing the arrays into train and test partitions. Setting np.random.seed when random_state is provided makes the split reproducible, matching sklearn's deterministic behavior. The function accepts **kwargs so callers can pass extra arguments without breaking — useful when patching code that uses stratify or shuffle. This mock is ideal for tests where you want to verify downstream logic (like model fitting or evaluation) without depending on sklearn's exact split details.

Common mistakes

  • Forgetting to pass `random_state` when reproducibility matters, causing flaky tests
  • Not handling list inputs — the mock assumes numpy arrays, so wrapping with `np.array` is needed first
  • Ignoring `shuffle=False` or `stratify` parameters that the real function supports

Variations

  1. Simplify further by returning a fixed split index pattern instead of random permutation for super-fast tests
  2. Use `unittest.mock.patch` to replace `sklearn.model_selection.train_test_split` in the module under test, delegating to the mock

Real-world use cases

  • Unit testing a classifier pipeline without installing scikit-learn in a minimal CI environment.
  • Validating feature engineering logic by ensuring train/test arrays align correctly with a custom split.
  • Writing fast test doubles for code that calls train_test_split so model training steps are isolated and repeatable.

Sponsored

Run locally

This sample needs third-party packages, so it cannot run in the browser IDE. Copy the code above, install the packages shown at the top, then run it in your own Python environment.

More from ML engineering pipelines

Keep learning

Related tutorials and quizzes for this topic.

Tutorials

Quizzes