Testing

pytest for pandas & numpy: Testing Data Transformations That Actually Work

HelpMeTest

HelpMeTest

— 5 min read
pytest for pandas & numpy: Testing Data Transformations That Actually Work

Data transformation bugs are silent killers. A shape mismatch, an off-by-one in a rolling window, a dtype coercion you didn't notice — these don't raise exceptions. They silently corrupt downstream models, dashboards, and decisions.

This guide focuses on pytest patterns specifically for pandas and numpy code: fixtures for test data, parametrize for edge cases, and assertion helpers that give you readable failures instead of cryptic stack traces.

Why Standard pytest Assertions Fall Short

The default assert df1 == df2 doesn't work for DataFrames — it raises an ambiguity error. And assert (df1 == df2).all().all() fails silently when shapes mismatch. You need specialized tools.

# DON'T do this
def test_transform():
    result = clean_dataframe(raw)
    assert result == expected  # ValueError: ambiguous truth value

# DO this
import pandas.testing as tm

def test_transform():
    result = clean_dataframe(raw)
    tm.assert_frame_equal(result, expected)

pandas.testing.assert_frame_equal gives you column-by-column diffs, dtype mismatches, and shape errors in one readable message.

Setting Up Fixtures for Reusable Test Data

Instead of defining DataFrames in every test, use fixtures:

# conftest.py
import pytest
import pandas as pd
import numpy as np

@pytest.fixture
def sample_sales_df():
    """Minimal sales DataFrame with known properties."""
    return pd.DataFrame({
        "date": pd.date_range("2024-01-01", periods=5),
        "product": ["A", "B", "A", "C", "B"],
        "revenue": [100.0, 200.0, 150.0, 75.0, 300.0],
        "quantity": [2, 4, 3, 1, 6],
    })

@pytest.fixture
def sales_with_nulls(sample_sales_df):
    """Sales DataFrame with injected nulls for null-handling tests."""
    df = sample_sales_df.copy()
    df.loc[1, "revenue"] = np.nan
    df.loc[3, "quantity"] = np.nan
    return df

Fixture composition (sales_with_nulls builds on sample_sales_df) keeps your test data DRY and clearly expresses what each test needs.

Testing DataFrame Transformations

Shape and Column Assertions

Always verify the structural contract of your transformation:

def test_add_revenue_per_unit(sample_sales_df):
    result = add_revenue_per_unit(sample_sales_df)
    
    # Shape preserved (rows unchanged)
    assert result.shape[0] == sample_sales_df.shape[0]
    
    # New column added
    assert "revenue_per_unit" in result.columns
    
    # Original columns preserved
    for col in sample_sales_df.columns:
        assert col in result.columns

Value-Level Assertions

For computed columns, test the math directly:

def test_revenue_per_unit_calculation(sample_sales_df):
    result = add_revenue_per_unit(sample_sales_df)
    
    # Vectorized check
    expected = sample_sales_df["revenue"] / sample_sales_df["quantity"]
    pd.testing.assert_series_equal(
        result["revenue_per_unit"],
        expected,
        check_names=False,  # ignore Series name differences
        rtol=1e-5,           # relative tolerance for floats
    )

Dtype Preservation

Type coercions are a common source of subtle bugs:

def test_dtypes_preserved(sample_sales_df):
    result = normalize_revenues(sample_sales_df)
    
    assert result["date"].dtype == "datetime64[ns]"
    assert result["product"].dtype == object  # string
    assert result["revenue"].dtype == np.float64

Parametrize for Edge Cases

Use @pytest.mark.parametrize to test the same logic against multiple scenarios without code duplication:

@pytest.mark.parametrize("revenue,quantity,expected", [
    (100.0, 2, 50.0),       # normal case
    (0.0, 5, 0.0),           # zero revenue
    (100.0, 1, 100.0),       # single unit
    (1e6, 1000, 1000.0),     # large values
])
def test_revenue_per_unit_values(revenue, quantity, expected):
    df = pd.DataFrame({"revenue": [revenue], "quantity": [quantity]})
    result = add_revenue_per_unit(df)
    assert abs(result["revenue_per_unit"].iloc[0] - expected) < 1e-9

This approach makes it trivial to add new test cases when you encounter a bug — just add a row to the parametrize list.

Testing numpy Array Transformations

For numpy-heavy code (feature engineering, signal processing, numerical methods):

import numpy as np

def normalize(arr: np.ndarray) -> np.ndarray:
    """Min-max normalize to [0, 1]."""
    return (arr - arr.min()) / (arr.max() - arr.min())

def test_normalize_range():
    arr = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
    result = normalize(arr)
    
    assert result.min() == pytest.approx(0.0)
    assert result.max() == pytest.approx(1.0)

def test_normalize_shape():
    arr = np.random.rand(100, 10)
    result = normalize(arr)
    assert result.shape == arr.shape

def test_normalize_dtype():
    arr = np.array([1, 2, 3], dtype=np.int32)
    result = normalize(arr)
    assert result.dtype == np.float64  # should upcast

pytest.approx handles float comparison correctly — it applies a relative tolerance of 1e-6 by default, which avoids false failures from floating-point arithmetic.

Testing with numpy.testing

For array comparisons, numpy.testing mirrors the pandas testing API:

import numpy.testing as npt

def test_rolling_mean():
    arr = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
    result = rolling_mean(arr, window=3)
    expected = np.array([np.nan, np.nan, 2.0, 3.0, 4.0])
    
    npt.assert_array_almost_equal(result, expected)
    # or for strict equality:
    npt.assert_array_equal(result[2:], expected[2:])

Handling NaN and Missing Data

Missing data tests deserve their own category — they reveal null-propagation bugs:

def test_aggregate_skips_nulls(sales_with_nulls):
    result = compute_total_revenue(sales_with_nulls)
    
    # Should sum non-null values only
    expected = sales_with_nulls["revenue"].sum()  # pandas sum() skips NaN
    assert result == pytest.approx(expected)

def test_null_rows_dropped(sales_with_nulls):
    result = drop_incomplete_rows(sales_with_nulls)
    
    # Original has 2 nulls, result should have 3 rows
    assert result.shape[0] == 3
    assert result.isnull().sum().sum() == 0

def test_null_rows_filled(sales_with_nulls):
    result = fill_nulls_with_median(sales_with_nulls)
    
    assert result.isnull().sum().sum() == 0
    # Median of [100, 150, 75, 300] = 125
    assert result["revenue"].median() == pytest.approx(125.0)

Groupby and Aggregation Testing

GroupBy operations have a common failure mode: you aggregate the wrong column or use the wrong function. Test the resulting shape AND values:

def test_revenue_by_product(sample_sales_df):
    result = total_revenue_by_product(sample_sales_df)
    
    # One row per unique product
    assert result.shape[0] == sample_sales_df["product"].nunique()
    
    # Correct totals
    assert result.loc["A", "revenue"] == pytest.approx(250.0)  # 100 + 150
    assert result.loc["B", "revenue"] == pytest.approx(500.0)  # 200 + 300
    assert result.loc["C", "revenue"] == pytest.approx(75.0)

Testing Time Series Transformations

Date-based logic (resampling, lag features, rolling windows) has many edge cases around period boundaries:

@pytest.fixture
def daily_timeseries():
    return pd.DataFrame({
        "date": pd.date_range("2024-01-01", periods=10, freq="D"),
        "value": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
    }).set_index("date")

def test_weekly_resample(daily_timeseries):
    result = resample_weekly(daily_timeseries)
    
    # 10 days → 2 full weeks
    assert result.shape[0] == 2
    # First week sum: 1+2+3+4+5+6+7 = 28
    assert result["value"].iloc[0] == 28

def test_lag_feature(daily_timeseries):
    result = add_lag(daily_timeseries, lag=1)
    
    # First row has NaN (no previous value)
    assert pd.isna(result["value_lag1"].iloc[0])
    # Second row has value of first row
    assert result["value_lag1"].iloc[1] == 1.0

Snapshot Testing for Complex Transformations

When a transformation is complex and you trust the current output, snapshot it:

# Save a snapshot during development:
# df.to_csv("tests/fixtures/expected_output.csv", index=False)

def test_full_pipeline_snapshot():
    result = full_feature_engineering_pipeline(raw_df)
    expected = pd.read_csv("tests/fixtures/expected_output.csv")
    
    pd.testing.assert_frame_equal(
        result.reset_index(drop=True),
        expected,
        check_dtype=False,  # allow int64/float64 flexibility
        atol=1e-4,
    )

Update the snapshot when you intentionally change the transformation, and treat unexpected changes as test failures.

Marking Slow Tests

Large DataFrame tests can be slow. Mark them so you can skip them during rapid iteration:

@pytest.mark.slow
def test_pipeline_on_1m_rows():
    large_df = generate_synthetic_data(n=1_000_000)
    result = full_pipeline(large_df)
    assert result.shape[0] == 1_000_000

# Run fast tests only:
# pytest -m "not slow"

# Run everything:
# pytest

Register the mark in pyproject.toml:

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

CI Configuration

Keep data transformation tests in your CI pipeline with a reasonable timeout:

# .github/workflows/test.yml
- name: Run data tests
  run: |
    pytest tests/data/ -v --timeout=60 -m "not slow"

Run slow tests nightly or on release branches:

- name: Run full data test suite
  if: github.ref == 'refs/heads/main'
  run: pytest tests/data/ -v --timeout=300

What to Actually Test

Focus your test budget on:

  1. Transformations that change shape — groupby, merge, pivot, resample
  2. Null handling — how does each step handle missing data?
  3. Type contracts — do outputs have the dtypes callers expect?
  4. Boundary conditions — empty DataFrames, single-row DataFrames, all-null columns
  5. Mathematical correctness — for derived features, verify the formula

Don't test pandas/numpy internals — they're already tested upstream. Test your code.

Conclusion

Solid pytest patterns for pandas and numpy are about three things: fixtures that create controlled test data, parametrize that covers edge cases without repetition, and specialized assertion helpers that give you readable failures. Apply these patterns consistently and you'll catch data transformation bugs at the source — before they corrupt your models or mislead your stakeholders.

Read more

Start now free