Data Pipeline Testing skill for Codex

A tested pipeline has fast unit tests for every transformation, a contract at each boundary where data changes hands, a golden-dataset check that catches any unintended change in outputs, and proof that running the same batch twice gives the same result. All of it runs in CI without production credentials.

Inspect first

1. Read AGENTS.md for the test command, orchestrator (Airflow, Dagster, Prefect, cron), engine (pandas, Polars, PySpark, SQL in a warehouse) and existing tests.
2. Map the pipeline: sources, each transformation step, sinks, and which steps are pure functions versus I/O. Untestable pipelines usually mix I/O and logic in one function.
3. Find past incidents (postmortems, issues): each is a test case waiting to be written.

Test layers

Layer	What it proves	Speed	Tools
Unit	One transformation is correct on hand-built rows	ms	pytest, pandas.testing, polars.testing
Contract	Data crossing a boundary has the agreed schema and invariants	ms-s	pandera, Pydantic, JSON Schema, dbt model contracts
Golden dataset	Whole pipeline output on fixed input is unchanged (or changes as intended)	s	Parquet/CSV snapshots in repo, DuckDB for diffs
Integration	Real I/O works: SQL dialect, warehouse writes, file formats	s-min	DuckDB, Testcontainers (Postgres, Kafka), local Spark
Orchestration	DAGs load, dependencies are wired, schedules valid	s	Airflow DagBag, Dagster Definitions validation
Runtime checks	Production data meets expectations each run	per run	hand off to $data-quality-checks

Spend most effort on unit and golden tests; they catch most regressions cheaply.

Make logic testable

Separate I/O from transformation so the middle is a pure function of DataFrames:

# pipeline/orders.py
import polars as pl

def extract(conn) -> pl.DataFrame: ...                      # I/O
def load(df: pl.DataFrame, conn, batch_date: str) -> None: ...  # I/O

def transform(orders: pl.DataFrame, fx: pl.DataFrame) -> pl.DataFrame:  # pure
    return (
        orders.filter(pl.col("status") != "test")
        .join(fx, on=["currency", "order_date"], how="left", validate="m:1")
        .with_columns((pl.col("amount") * pl.col("rate_to_eur")).round(2).alias("amount_eur"))
        .unique(subset="order_id", keep="last", maintain_order=True)
    )

validate="m:1" makes Polars raise if the FX table has duplicate keys; pandas has merge(..., validate="many_to_one").

Unit tests with tiny fixtures

# tests/test_orders_transform.py
from datetime import date
import polars as pl
from polars.testing import assert_frame_equal
import pytest
from pipeline.orders import transform

@pytest.fixture
def fx():
    return pl.DataFrame({"currency": ["USD", "EUR"], "order_date": [date(2025, 10, 1)] * 2,
                         "rate_to_eur": [0.92, 1.0]})

def make_orders(**overrides):
    base = {"order_id": ["o1"], "status": ["paid"], "currency": ["USD"],
            "order_date": [date(2025, 10, 1)], "amount": [100.0]}
    return pl.DataFrame({**base, **overrides})

def test_converts_to_eur(fx):
    out = transform(make_orders(), fx)
    assert out["amount_eur"].to_list() == [92.0]

def test_drops_test_orders(fx):
    assert transform(make_orders(status=["test"]), fx).is_empty()

def test_duplicate_order_keeps_last(fx):
    orders = make_orders(order_id=["o1", "o1"], status=["paid", "paid"], currency=["EUR", "EUR"],
                         order_date=[date(2025, 10, 1)] * 2, amount=[10.0, 12.0])
    assert transform(orders, fx)["amount_eur"].to_list() == [12.0]

def test_missing_fx_rate_yields_null_not_crash(fx):
    out = transform(make_orders(currency=["JPY"]), fx)
    assert out["amount_eur"].null_count() == 1

def test_duplicate_fx_rates_raise(fx):
    with pytest.raises(Exception):
        transform(make_orders(), pl.concat([fx, fx]))

Edge-case checklist for fixtures

- [ ] Empty input (zero rows, but correct schema)
- [ ] Nulls in every nullable column, including join keys
- [ ] Duplicates on the business key
- [ ] Late-arriving rows for an earlier batch date
- [ ] Time-zone boundaries: 23:30 local on the last day of a month, DST change days
- [ ] Unicode and very long strings; leading zeros in IDs
- [ ] Extreme numbers: 0, negative, very large, NaN vs null
- [ ] Schema drift: an extra column, a missing optional column

Contracts at boundaries

import polars as pl
import pandera.polars as pa

class OrdersOut(pa.DataFrameModel):
    order_id: str = pa.Field(unique=True)
    amount_eur: float = pa.Field(ge=0, nullable=True)
    order_date: pl.Date

    class Config:
        strict = True          # unexpected columns fail

def transform_checked(orders, fx):
    return OrdersOut.validate(transform(orders, fx))

Validate on the producer side before writing, and on the consumer side for data you do not control. For dbt models, use model contracts and unit tests; hand off to $dbt-testing.

Golden datasets

Freeze a small, representative input (hundreds to a few thousand rows, anonymised) and the expected output in the repo:

# tests/test_golden.py
from pathlib import Path
import polars as pl
from polars.testing import assert_frame_equal
from pipeline.orders import transform

GOLDEN = Path(__file__).parent / "golden"

def test_orders_golden(request):
    out = transform(pl.read_parquet(GOLDEN / "orders_in.parquet"), pl.read_parquet(GOLDEN / "fx.parquet"))
    out = out.sort("order_id")
    expected_path = GOLDEN / "orders_out.parquet"
    if request.config.getoption("--update-golden", default=False):
        out.write_parquet(expected_path)
    assert_frame_equal(out, pl.read_parquet(expected_path), check_dtypes=True, abs_tol=1e-9)

# tests/conftest.py
def pytest_addoption(parser):
    parser.addoption("--update-golden", action="store_true", help="rewrite golden outputs")

• Sort deterministically before comparing; row order is not a contract.
• Updating goldens is a reviewed change: the PR must show and explain the diff. Summarise it with DuckDB: select * from 'new.parquet' except select * from 'old.parquet'.
• Never update goldens just to make a red test green without understanding why.

Idempotency and re-runs

def test_rerun_same_batch_is_idempotent(duckdb_conn):
    run_pipeline(duckdb_conn, batch_date="2025-10-01")
    first = duckdb_conn.sql("select count(*), sum(amount_eur) from fct_orders").fetchone()
    run_pipeline(duckdb_conn, batch_date="2025-10-01")
    assert duckdb_conn.sql("select count(*), sum(amount_eur) from fct_orders").fetchone() == first

Also test: a backfill of an older date does not touch other partitions, and a failed run midway leaves the target unchanged (write to staging, then swap).

Integration and orchestration

• SQL that targets Postgres: run it against Postgres in Testcontainers (testcontainers.postgres.PostgresContainer("postgres:17")), which needs Docker and may need network to pull images. DuckDB is a fast stand-in only when the SQL is portable.
• Airflow: a test that every DAG imports without errors and has no cycles:

from airflow.models import DagBag

def test_dags_import():
    bag = DagBag(dag_folder="dags/", include_examples=False)
    assert not bag.import_errors, bag.import_errors

• Dagster: load Definitions in a test and materialise assets with in-memory I/O managers.
• Prefect: call flow functions directly in tests; they run locally without a server.

Gotchas

• Tests that read production or shared staging data are flaky and leak data into CI. Use synthetic or anonymised fixtures only.
• datetime.now() inside transforms makes outputs untestable; pass the batch date in.
• Float comparisons on money: use tolerances or decimal types consistently.
• Spark tests: one shared local SparkSession fixture with scope="session" keeps the suite fast; set spark.sql.shuffle.partitions to a small number.
• Mocking the database instead of running SQL hides dialect bugs; test SQL on a real engine.

Verification

pytest -q tests/                       # unit + golden + contract, should run in seconds
pytest -q -m integration tests/        # if marked; needs Docker

• Mutate one transformation deliberately (change a filter); at least one test must fail.
• CI runs the suite on every PR without production secrets.

Delivery template

**Pipeline:** orders_daily (Polars, Airflow)
**Refactor:** transform() split from I/O; batch_date injected
**Tests added:** 14 unit (edge-case checklist), 1 golden (2,400 rows), 2 contract, 1 idempotency (DuckDB), DAG import
**Runtime:** 3.1 s locally; integration (Postgres via Testcontainers) 22 s, marked `integration`
**Mutation check:** removing the test-order filter fails 2 tests
**Golden diff:** none (refactor is output-identical)