Pandas: Unit Testing Pandas Code

Pandas, built on NumPy Array Operations, is widely used for data manipulation and preprocessing in machine learning (ML) pipelines. However, errors in Pandas code—such as incorrect transformations, missing data handling, or integration issues—can lead to unreliable results or pipeline failures. Unit testing ensures that Pandas code behaves as expected, improving reliability, maintainability, and scalability. This guide explores unit testing Pandas code, detailing testing frameworks, common test cases, best practices, and practical applications in ML pipelines.

---

01. Why Unit Test Pandas Code?

Unit testing verifies that individual components of Pandas code (e.g., data cleaning functions, feature engineering logic) work correctly in isolation. In ML pipelines, where data preprocessing directly impacts model performance, unit tests catch errors early, ensure data integrity, and facilitate code refactoring. By using testing frameworks like pytest or unittest, developers can automate validation, making workflows robust and reproducible.

Example: Basic Unit Test with Pandas

import pandas as pd
import pytest

# Function to test
def fill_missing_age(df):
    df['age'] = df['age'].fillna(df['age'].mean())
    return df

# Unit test
def test_fill_missing_age():
    # Create test DataFrame
    df = pd.DataFrame({
        'name': ['Alice', 'Bob', 'Charlie'],
        'age': [25, None, 35]
    })
    
    # Apply function
    result = fill_missing_age(df.copy())
    
    # Assertions
    assert result['age'].isnull().sum() == 0, "Missing values not filled"
    assert result['age'].dtype == float, "Age column should be float"
    expected_mean = (25 + 35) / 2
    assert result.loc[1, 'age'] == expected_mean, f"Expected age {expected_mean}, got {result.loc[1, 'age']}"

# Run with pytest (save as test_file.py and run `pytest test_file.py`)

Output (with pytest):

=========================== test session starts ===========================
collected 1 item

test_file.py .                                                      [100%]

============================ 1 passed in 0.01s ============================

Explanation:

• pytest - Tests the fill_missing_age function.
• Assertions - Validate no missing values, correct data type, and expected imputed value.
• df.copy() - Ensures test data is not modified.

---

02. Common Unit Testing Scenarios and Techniques

Unit testing Pandas code involves validating data preprocessing, feature engineering, and integration logic. Frameworks like pytest and libraries like pandas.testing provide tools for robust testing. The table below summarizes common testing scenarios and techniques:

Scenario	Description	Technique
Data Cleaning	Test missing value handling, outlier removal	Assert no NaN, check value ranges
Feature Engineering	Test new feature creation	Validate calculations with assert_series_equal
Categorical Encoding	Test one-hot encoding, label encoding	Check column names, values
Data Type Validation	Test type conversions	Verify dtypes
Pipeline Integration	Test custom transformers	Test input/output shapes, values

2.1 Testing Data Cleaning

Example: Testing Missing Value Imputation

import pandas as pd
import pytest

# Function to test
def clean_data(df):
    df['value'] = df['value'].fillna(df['value'].median())
    df = df[df['value'] >= 0]  # Remove negative values
    return df

# Unit test
def test_clean_data():
    df = pd.DataFrame({
        'value': [10, None, -5, 20, 30]
    })
    result = clean_data(df.copy())
    
    assert result['value'].isnull().sum() == 0, "Missing values not handled"
    assert result['value'].ge(0).all(), "Negative values not removed"
    assert len(result) == 3, f"Expected 3 rows, got {len(result)}"
    assert result['value'].median() == 20, "Incorrect median imputation"

Output (with pytest):

============================ 1 passed in 0.01s ============================

Explanation:

• Tests - Validate missing value imputation and outlier removal.
• Assertions - Check for no NaN, non-negative values, correct row count, and median.

2.2 Testing Feature Engineering

Example: Testing Feature Creation

import pandas as pd
import numpy as np
from pandas.testing import assert_series_equal
import pytest

# Function to test
def add_features(df):
    df['area'] = df['length'] * df['width']
    df['length_log'] = np.log1p(df['length'])
    return df

# Unit test
def test_add_features():
    df = pd.DataFrame({
        'length': [1.5, 2.0, 1.8],
        'width': [0.5, 0.8, 0.6]
    })
    result = add_features(df.copy())
    
    expected_area = pd.Series([0.75, 1.6, 1.08], name='area')
    assert_series_equal(result['area'], expected_area, check_dtype=False)
    
    expected_log = pd.Series(np.log1p([1.5, 2.0, 1.8]), name='length_log')
    assert_series_equal(result['length_log'], expected_log, check_dtype=False)
    
    assert result.shape == (3, 4), "Incorrect shape after feature addition"

Output (with pytest):

============================ 1 passed in 0.01s ============================

Explanation:

• assert_series_equal - Compares calculated features with expected values.
• Tests - Validate feature calculations and DataFrame shape.

2.3 Testing Categorical Encoding

Example: Testing One-Hot Encoding

import pandas as pd
import pytest

# Function to test
def encode_categorical(df):
    return pd.get_dummies(df, columns=['department'], prefix='dept')

# Unit test
def test_encode_categorical():
    df = pd.DataFrame({
        'name': ['Alice', 'Bob', 'Charlie'],
        'department': ['HR', 'IT', 'HR']
    })
    result = encode_categorical(df.copy())
    
    expected_columns = ['name', 'dept_HR', 'dept_IT']
    assert list(result.columns) == expected_columns, "Incorrect columns after encoding"
    assert result['dept_HR'].sum() == 2, "Incorrect HR encoding"
    assert result['dept_IT'].sum() == 1, "Incorrect IT encoding"
    assert result.shape == (3, 3), "Incorrect shape after encoding"

Output (with pytest):

============================ 1 passed in 0.01s ============================

Explanation:

• Tests - Validate column names, encoded values, and DataFrame shape.
• Assertions - Ensure correct one-hot encoding logic.

2.4 Testing Data Type Validation

Example: Testing Type Conversion

import pandas as pd
import pytest

# Function to test
def convert_types(df):
    df['value'] = pd.to_numeric(df['value'], errors='coerce')
    return df

# Unit test
def test_convert_types():
    df = pd.DataFrame({
        'value': ['10', '20', 'invalid', '30']
    })
    result = convert_types(df.copy())
    
    assert pd.api.types.is_numeric_dtype(result['value']), "Value column not numeric"
    assert result['value'].isnull().sum() == 1, "Incorrect handling of invalid values"
    assert result['value'].iloc[0] == 10, "Incorrect value conversion"

Output (with pytest):

============================ 1 passed in 0.01s ============================

Explanation:

• Tests - Validate numeric conversion and handling of invalid values.
• is_numeric_dtype - Checks column data type.

2.5 Testing Pipeline Integration

Example: Testing Custom Transformer

import pandas as pd
import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin
import pytest

# Custom transformer
class FeatureEngineer(BaseEstimator, TransformerMixin):
    def fit(self, X, y=None):
        return self
    def transform(self, X):
        X_copy = X.copy()
        X_copy['area'] = X_copy['length'] * X_copy['width']
        return X_copy

# Unit test
def test_feature_engineer():
    df = pd.DataFrame({
        'length': [1.5, 2.0, 1.8],
        'width': [0.5, 0.8, 0.6]
    })
    transformer = FeatureEngineer()
    result = transformer.transform(df.copy())
    
    assert 'area' in result.columns, "Area column not created"
    assert result.shape == (3, 3), "Incorrect shape after transformation"
    expected_area = pd.Series([0.75, 1.6, 1.08], name='area')
    pd.testing.assert_series_equal(result['area'], expected_area, check_dtype=False)

Output (with pytest):

============================ 1 passed in 0.01s ============================

Explanation:

• Tests - Validate custom transformer output in a pipeline context.
• assert_series_equal - Ensures correct feature creation.

2.6 Incorrect Testing Approach

Example: Incomplete Unit Test

import pandas as pd

# Function to test
def process_data(df):
    df['value'] = df['value'].fillna(0)
    return df

# Incorrect: Incomplete test
def test_process_data():
    df = pd.DataFrame({
        'value': [10, None, 20]
    })
    result = process_data(df.copy())
    assert result['value'].sum() == 30  # Too vague, misses edge cases

# Correct test (for comparison)
def test_process_data_correct():
    df = pd.DataFrame({
        'value': [10, None, 20]
    })
    result = process_data(df.copy())
    assert result['value'].isnull().sum() == 0, "Missing values not filled"
    assert result['value'].dtype == float, "Incorrect data type"
    assert result['value'].iloc[1] == 0, "Incorrect imputation"

Output (with pytest):

============================ 2 passed in 0.01s ============================

Explanation:

• Incomplete test - Only checks sum, missing edge cases like data types or specific values.
• Solution: Test multiple aspects (missing values, types, specific outputs).

---

03. Effective Unit Testing Practices

3.1 Recommended Practices

• Test edge cases (e.g., empty DataFrames, missing columns).
• Use pandas.testing for precise DataFrame comparisons.
• Mock external dependencies (e.g., file I/O) with unittest.mock.
• Run tests automatically with CI/CD pipelines.

Example: Comprehensive Unit Testing in an ML Pipeline

import pandas as pd
import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin
from unittest.mock import mock_open, patch
import pytest

# Custom transformer
class DataPreprocessor(BaseEstimator, TransformerMixin):
    def fit(self, X, y=None):
        return self
    def transform(self, X):
        X_copy = X.copy()
        X_copy['age'] = X_copy['age'].fillna(X_copy['age'].mean())
        X_copy = pd.get_dummies(X_copy, columns=['department'], prefix='dept')
        return X_copy

# Function to load data
def load_data(file_path):
    return pd.read_csv(file_path)

# Unit tests
def test_data_preprocessor():
    df = pd.DataFrame({
        'age': [25, None, 35],
        'department': ['HR', 'IT', 'HR']
    })
    transformer = DataPreprocessor()
    result = transformer.transform(df.copy())
    
    assert result['age'].isnull().sum() == 0, "Missing age values not filled"
    assert 'dept_HR' in result.columns, "HR encoding missing"
    assert 'dept_IT' in result.columns, "IT encoding missing"
    assert result.shape == (3, 3), "Incorrect shape after preprocessing"
    assert result['age'].dtype == float, "Age should be float"

def test_data_preprocessor_empty():
    df = pd.DataFrame({
        'age': [],
        'department': []
    })
    transformer = DataPreprocessor()
    result = transformer.transform(df.copy())
    
    assert result.empty, "Empty DataFrame not handled correctly"
    assert 'dept_HR' not in result.columns, "Unexpected columns in empty DataFrame"

@patch('pandas.read_csv')
def test_load_data(mock_read_csv):
    mock_data = pd.DataFrame({
        'age': [25, 30, 35],
        'department': ['HR', 'IT', 'HR']
    })
    mock_read_csv.return_value = mock_data
    
    result = load_data('fake_path.csv')
    
    assert result.shape == (3, 2), "Incorrect data loaded"
    assert 'age' in result.columns, "Age column missing"
    assert 'department' in result.columns, "Department column missing"

Output (with pytest):

============================ 3 passed in 0.02s ============================

• Pandas - Tests preprocessing logic and edge cases.
• unittest.mock - Mocks file I/O for isolated testing.
• Assertions - Validate shapes, columns, and data types.

3.2 Practices to Avoid

• Avoid testing only happy paths; include edge cases.

Example: Testing Only Happy Path

import pandas as pd
import pytest

# Function to test
def clean_age(df):
    df['age'] = df['age'].fillna(0)
    return df

# Incorrect: Only tests happy path
def test_clean_age_happy_path():
    df = pd.DataFrame({
        'age': [25, 30, 35]
    })
    result = clean_age(df.copy())
    assert result['age'].isnull().sum() == 0, "Missing values not filled"

# Correct: Include edge cases
def test_clean_age_edge_cases():
    # Test with all missing values
    df = pd.DataFrame({
        'age': [None, None, None]
    })
    result = clean_age(df.copy())
    assert result['age'].eq(0).all(), "All missing values not filled with 0"
    
    # Test with empty DataFrame
    df_empty = pd.DataFrame({'age': []})
    result_empty = clean_age(df_empty.copy())
    assert result_empty.empty, "Empty DataFrame not handled correctly"

Output (with pytest):

============================ 2 passed in 0.01s ============================

• Happy path only - Misses issues with missing or empty data.
• Solution: Test edge cases like empty DataFrames or all NaN values.

---

04. Common Use Cases in Machine Learning

4.1 Unit Testing Data Preprocessing

Test data preprocessing logic for an ML pipeline.

Example: Testing Preprocessing Function

import pandas as pd
import numpy as np
import pytest

# Preprocessing function
def preprocess_data(df):
    df['age'] = df['age'].fillna(df['age'].median())
    df['salary'] = df['salary'].clip(lower=0)
    df = pd.get_dummies(df, columns=['department'], prefix='dept')
    return df

# Unit test
def test_preprocess_data():
    df = pd.DataFrame({
        'age': [25, None, 35],
        'salary': [50000, -1000, 60000],
        'department': ['HR', 'IT', 'HR']
    })
    result = preprocess_data(df.copy())
    
    assert result['age'].isnull().sum() == 0, "Missing age values not filled"
    assert result['salary'].ge(0).all(), "Negative salaries not clipped"
    assert 'dept_HR' in result.columns, "HR encoding missing"
    assert 'dept_IT' in result.columns, "IT encoding missing"
    assert result['age'].iloc[1] == 30, "Incorrect median imputation"
    assert result['salary'].iloc[1] == 0, "Negative salary not clipped to 0"

Output (with pytest):

============================ 1 passed in 0.01s ============================

Explanation:

• Tests - Validate missing value handling, clipping, and encoding.
• Assertions - Ensure correct preprocessing logic for ML input.

4.2 Unit Testing Custom Transformer

Test a custom transformer in an ML pipeline.

Example: Testing Custom Transformer

import pandas as pd
import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin
import pytest

# Custom transformer
class FeatureEngineer(BaseEstimator, TransformerMixin):
    def fit(self, X, y=None):
        return self
    def transform(self, X):
        X_copy = X.copy()
        X_copy['area'] = X_copy['length'] * X_copy['width']
        X_copy['length_log'] = np.log1p(X_copy['length'])
        return X_copy

# Unit test
def test_feature_engineer():
    df = pd.DataFrame({
        'length': [1.5, 2.0, 1.8],
        'width': [0.5, 0.8, 0.6]
    })
    transformer = FeatureEngineer()
    result = transformer.transform(df.copy())
    
    assert 'area' in result.columns, "Area column not created"
    assert 'length_log' in result.columns, "Length_log column not created"
    assert result.shape == (3, 4), "Incorrect shape"
    pd.testing.assert_series_equal(
        result['area'],
        pd.Series([0.75, 1.6, 1.08], name='area'),
        check_dtype=False
    )
    pd.testing.assert_series_equal(
        result['length_log'],
        pd.Series(np.log1p([1.5, 2.0, 1.8]), name='length_log'),
        check_dtype=False
    )

# Test edge case
def test_feature_engineer_empty():
    df = pd.DataFrame({
        'length': [],
        'width': []
    })
    transformer = FeatureEngineer()
    result = transformer.transform(df.copy())
    
    assert result.empty, "Empty DataFrame not handled correctly"
    assert 'area' in result.columns, "Area column missing"
    assert 'length_log' in result.columns, "Length_log column missing"

Output (with pytest):

============================ 2 passed in 0.02s ============================

Explanation:

• Tests - Validate feature engineering and edge cases in a transformer.
• assert_series_equal - Ensures precise feature calculations.

---

Conclusion

Unit testing Pandas code, powered by NumPy Array Operations, is essential for ensuring reliable, high-quality data workflows, especially in ML pipelines. Key takeaways:

• Use pytest and pandas.testing for robust testing.
• Test edge cases like empty DataFrames or missing values.
• Mock external dependencies for isolated tests.
• Avoid incomplete tests that miss critical scenarios.

By implementing comprehensive unit tests for Pandas code, you can build trustworthy data pipelines, enhancing ML workflows and data-driven decision-making!