pytest for Data Pipelines: Fixtures, Parametrize, and Integration Testing

pytest for Data Pipelines: Fixtures, Parametrize, and Integration Testing

Data pipelines fail in ways that pure unit tests don't catch. A transformation that works on clean test data breaks on real-world schemas with extra columns, unexpected nulls, or timezone-aware timestamps. A join that looks correct drops 40% of rows silently.

This guide covers pytest patterns specifically for ETL and data pipeline code: how to structure fixtures, test pipeline stages in isolation, test the full pipeline against a real (but minimal) database, and catch the class of bugs that only emerge at integration boundaries.

The Pipeline Testing Stack

For Python data pipelines, you'll typically use:

  • pytest — test runner and fixture system
  • pytest-mock — for mocking external systems (S3, APIs, databases)
  • SQLite or DuckDB — in-memory databases for integration tests
  • pandas.testing — DataFrame assertion utilities
  • factory_boy or pytest fixtures — test data factories

Install them:

pip install pytest pytest-mock duckdb pandas factory-boy

Structuring Your Test Suite

Organize tests to mirror the pipeline stages:

tests/
  conftest.py           # Shared fixtures
  unit/
    test_extract.py     # Individual transform functions
    test_transform.py
    test_load.py
  integration/
    test_pipeline.py    # Full pipeline against in-memory DB
  fixtures/
    raw_events.csv      # Static test data files
    expected_output.csv

The conftest.py Strategy

conftest.py is where pytest fixtures live. Structure it by concern:

# tests/conftest.py
import pytest
import duckdb
import pandas as pd
from pathlib import Path

FIXTURES_DIR = Path(__file__).parent / "fixtures"

# ── Database fixtures ──────────────────────────────────────────

@pytest.fixture(scope="session")
def duckdb_conn():
    """In-memory DuckDB connection, shared across the test session."""
    conn = duckdb.connect(":memory:")
    yield conn
    conn.close()

@pytest.fixture
def fresh_db(duckdb_conn):
    """Reset tables before each test."""
    duckdb_conn.execute("DROP TABLE IF EXISTS events")
    duckdb_conn.execute("DROP TABLE IF EXISTS users")
    yield duckdb_conn
    # Cleanup happens on next invocation

# ── Data fixtures ──────────────────────────────────────────────

@pytest.fixture
def raw_events():
    """Minimal raw events DataFrame — known, controlled data."""
    return pd.DataFrame({
        "event_id": [1, 2, 3, 4, 5],
        "user_id": [101, 102, 101, 103, 102],
        "event_type": ["click", "view", "purchase", "click", "view"],
        "timestamp": pd.to_datetime([
            "2024-01-01 10:00:00",
            "2024-01-01 10:05:00",
            "2024-01-01 10:10:00",
            "2024-01-01 10:15:00",
            "2024-01-01 10:20:00",
        ]),
        "amount": [None, None, 49.99, None, None],
    })

@pytest.fixture
def raw_users():
    return pd.DataFrame({
        "user_id": [101, 102, 103],
        "email": ["alice@example.com", "bob@example.com", "carol@example.com"],
        "signup_date": pd.to_datetime(["2023-12-01", "2023-11-15", "2024-01-01"]),
        "plan": ["pro", "free", "pro"],
    })

Testing Extract Functions

Extract functions (reading from S3, databases, APIs) should be tested with mocks to avoid network calls in CI:

# tests/unit/test_extract.py
from unittest.mock import patch, MagicMock
import pandas as pd
import pytest
from my_pipeline.extract import read_events_from_s3

def test_read_events_returns_dataframe(mocker):
    """read_events_from_s3 should return a DataFrame with the right columns."""
    mock_df = pd.DataFrame({
        "event_id": [1, 2],
        "user_id": [101, 102],
        "event_type": ["click", "view"],
    })
    mocker.patch("my_pipeline.extract.pd.read_parquet", return_value=mock_df)
    
    result = read_events_from_s3("s3://bucket/path/events.parquet")
    
    assert isinstance(result, pd.DataFrame)
    assert list(result.columns) == ["event_id", "user_id", "event_type"]

def test_read_events_raises_on_missing_columns(mocker):
    """Should raise ValueError when required columns are missing."""
    incomplete_df = pd.DataFrame({"event_id": [1]})  # missing user_id, event_type
    mocker.patch("my_pipeline.extract.pd.read_parquet", return_value=incomplete_df)
    
    with pytest.raises(ValueError, match="Missing required columns"):
        read_events_from_s3("s3://bucket/path/events.parquet")

Testing Transform Functions

Transform functions should be pure (input → output, no side effects), making them easy to unit test:

# tests/unit/test_transform.py
import pandas as pd
import pytest
from my_pipeline.transform import (
    filter_purchase_events,
    add_user_tenure_days,
    normalize_event_types,
)

def test_filter_purchase_events(raw_events):
    result = filter_purchase_events(raw_events)
    
    assert result.shape[0] == 1  # Only 1 purchase in fixture
    assert (result["event_type"] == "purchase").all()
    assert result["amount"].notna().all()  # Purchases have amounts

def test_filter_purchase_events_empty_input():
    """Edge case: empty DataFrame should return empty DataFrame."""
    empty = pd.DataFrame(columns=["event_id", "user_id", "event_type", "amount"])
    result = filter_purchase_events(empty)
    assert result.shape[0] == 0
    assert list(result.columns) == list(empty.columns)

def test_add_user_tenure_days(raw_events, raw_users):
    result = add_user_tenure_days(raw_events, raw_users, reference_date="2024-01-01")
    
    assert "tenure_days" in result.columns
    # User 101 signed up 2023-12-01 → 31 days tenure
    user_101_events = result[result["user_id"] == 101]
    assert (user_101_events["tenure_days"] == 31).all()

@pytest.mark.parametrize("raw,expected", [
    ("click", "CLICK"),
    ("VIEW", "VIEW"),
    ("Purchase", "PURCHASE"),
    ("add_to_cart", "ADD_TO_CART"),
])
def test_normalize_event_types(raw, expected):
    df = pd.DataFrame({"event_type": [raw]})
    result = normalize_event_types(df)
    assert result["event_type"].iloc[0] == expected

Testing Joins (The Dangerous Part)

Joins are where silent data loss happens. Always test:

  1. Output row count vs. input row count
  2. Whether join keys exist in both sides
  3. Null handling in join keys
def test_enrich_events_with_users(raw_events, raw_users):
    from my_pipeline.transform import enrich_events_with_users
    
    result = enrich_events_with_users(raw_events, raw_users)
    
    # No rows should be lost — all event user_ids exist in users
    assert result.shape[0] == raw_events.shape[0]
    
    # User columns added
    assert "email" in result.columns
    assert "plan" in result.columns
    
    # Values correct
    alice_events = result[result["user_id"] == 101]
    assert (alice_events["email"] == "alice@example.com").all()

def test_enrich_events_unknown_user_id():
    """Events with unknown user_ids should be handled (not silently dropped)."""
    events = pd.DataFrame({
        "event_id": [1],
        "user_id": [999],  # Does not exist in users
        "event_type": ["click"],
    })
    users = pd.DataFrame({"user_id": [101], "email": ["alice@example.com"]})
    
    from my_pipeline.transform import enrich_events_with_users
    result = enrich_events_with_users(events, users)
    
    # Verify the implementation's contract: left join keeps the event
    assert result.shape[0] == 1
    assert pd.isna(result["email"].iloc[0])

Integration Testing with DuckDB

Test the full pipeline against an in-memory database:

# tests/integration/test_pipeline.py
import pandas as pd
import duckdb
import pytest
from my_pipeline.pipeline import run_events_pipeline

def test_full_pipeline_produces_correct_output(fresh_db, raw_events, raw_users):
    """End-to-end: raw events + users → aggregated purchase summary."""
    
    # Load test data into the in-memory DB
    fresh_db.register("events_view", raw_events)
    fresh_db.register("users_view", raw_users)
    
    # Run the pipeline
    result = run_events_pipeline(conn=fresh_db)
    
    # Verify shape: one row per user with purchases
    assert result.shape[0] == 1  # Only user 101 made a purchase
    
    # Verify values
    row = result.iloc[0]
    assert row["user_id"] == 101
    assert row["total_purchases"] == 1
    assert abs(row["total_revenue"] - 49.99) < 0.01
    assert row["plan"] == "pro"

def test_pipeline_handles_no_purchases(fresh_db, raw_users):
    """Pipeline with zero purchase events should return empty DataFrame."""
    no_purchases = pd.DataFrame({
        "event_id": [1, 2],
        "user_id": [101, 102],
        "event_type": ["click", "view"],
        "timestamp": pd.to_datetime(["2024-01-01", "2024-01-01"]),
        "amount": [None, None],
    })
    
    fresh_db.register("events_view", no_purchases)
    fresh_db.register("users_view", raw_users)
    
    result = run_events_pipeline(conn=fresh_db)
    assert result.shape[0] == 0

Testing Load Functions

Load functions (writing to databases, S3, data warehouses) should be tested against the in-memory DB and verified by reading back:

def test_load_results_to_db(fresh_db):
    from my_pipeline.load import load_purchase_summary
    
    summary = pd.DataFrame({
        "user_id": [101, 102],
        "total_purchases": [3, 1],
        "total_revenue": [149.97, 29.99],
    })
    
    load_purchase_summary(summary, conn=fresh_db)
    
    # Verify it landed correctly
    result = fresh_db.execute("SELECT * FROM purchase_summary ORDER BY user_id").df()
    
    assert result.shape[0] == 2
    assert abs(result["total_revenue"].sum() - 179.96) < 0.01

Idempotency Testing

Pipelines should be safe to re-run. Test that running twice produces the same result:

def test_pipeline_is_idempotent(fresh_db, raw_events, raw_users):
    fresh_db.register("events_view", raw_events)
    fresh_db.register("users_view", raw_users)
    
    result_1 = run_events_pipeline(conn=fresh_db)
    result_2 = run_events_pipeline(conn=fresh_db)
    
    pd.testing.assert_frame_equal(
        result_1.reset_index(drop=True),
        result_2.reset_index(drop=True),
    )

Parametrize for Multiple Scenarios

Use parametrize to run the same pipeline test with different input shapes:

@pytest.mark.parametrize("n_events,n_purchases_expected", [
    (0, 0),
    (1, 1),
    (100, 23),  # known ratio in synthetic data generator
])
def test_pipeline_purchase_count(n_events, n_purchases_expected, fresh_db, raw_users):
    events = generate_synthetic_events(n=n_events, purchase_ratio=0.23)
    fresh_db.register("events_view", events)
    fresh_db.register("users_view", raw_users)
    
    result = run_events_pipeline(conn=fresh_db)
    assert result["total_purchases"].sum() == n_purchases_expected

CI Integration

Structure your pytest configuration to run unit tests fast and integration tests separately:

# pyproject.toml
[tool.pytest.ini_options]
testpaths = ["tests"]
markers = [
    "unit: fast unit tests (no I/O)",
    "integration: slower integration tests with in-memory DB",
    "slow: very slow tests, skip on PR",
]

[tool.pytest.ini_options.filterwarnings]
ignore = "DeprecationWarning"
# .github/workflows/test.yml
jobs:
  unit-tests:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with: { python-version: "3.12" }
      - run: pip install -r requirements-dev.txt
      - run: pytest tests/unit/ -v --tb=short

  integration-tests:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with: { python-version: "3.12" }
      - run: pip install -r requirements-dev.txt
      - run: pytest tests/integration/ -v --timeout=120

What to Always Test in Data Pipelines

Concern Test Pattern
Row count preservation Assert input vs. output shape
Silent data loss in joins Test unknown key handling
Null propagation Fixtures with injected nulls
Dtype contracts Assert specific dtypes on output
Idempotency Run pipeline twice, compare results
Empty input Pass zero-row DataFrames
Large scale (smoke) 1M row smoke test, mark slow

Conclusion

pytest's fixture system and parametrize make it practical to write comprehensive data pipeline tests without massive amounts of boilerplate. The key insight: test each stage in isolation with unit tests, then wire it all together with a fast in-memory database integration test. Catch the bugs at the stage boundary, not in production.

Read more

Start now free