Testing Data Pipelines Like You Mean It: A pytest Crash Course for Data Engineers

Most data engineers write pipelines the way most people write shell scripts: run it, eyeball the output, ship it. That works right up until a schema changes upstream, a null slips through a join, or someone "fixes" a transformation and silently breaks three downstream tables. By then the bug isn't your problem anymore — it's a bad number in someone's dashboard.

Software engineers solved this problem decades ago with automated testing. Data engineering has been slower to adopt the habit, partly because our code touches messy external reality (files, databases, clusters) in a way a typical web app doesn't. But that's exactly why testing matters more here, not less. This article is a practical, DE-flavored crash course in pytest — the dominant Python testing framework — plus the patterns you actually need for pandas, Polars, and PySpark pipelines.

Why bother testing a data pipeline?

A few concrete failure modes that tests catch before production does:

• A column gets renamed upstream and your join silently produces all-null matches instead of erroring.
• A "cleaning" function that's supposed to drop duplicates accidentally drops valid rows too.
• A date-parsing function works on your local machine's locale and breaks in the CI environment.
• A refactor changes an aggregation from sum to mean and nobody notices until finance asks why revenue looks 90% smaller.

None of these require exotic testing techniques. They require the habit of writing small, deterministic checks against small, deterministic inputs — which is exactly what pytest is built for.

Where pytest fits — and where it doesn't

Before diving in, it's worth being precise about scope, because "testing a data pipeline" actually covers two different questions, and conflating them is a common source of confusion:

• Is my code correct? Given a known input, does the transformation logic produce the right output? This is a property of your code, and it doesn't change based on what day it is or what a source system decided to send you.
• Is today's data correct? Even with perfect code, a source system can start sending nulls, a partner feed can drop 90% of its rows overnight, a foreign key can stop resolving. This is a property of the data currently flowing through the system, and no amount of code testing can catch it, because the code was never wrong.

pytest answers the first question. Tools like dbt test, Great Expectations, and Soda answer the second. They're not competitors — they run at different times, against different inputs, and catch different bugs:

A useful rule of thumb when you're not sure which bucket a check belongs in: if the same check would fail identically on a completely different day's data, it's a code test; if it depends on what actually arrived today, it's a data test. "Does calculate_discount cap at 50%?" is always true or always false regardless of the date — code test. "Did today's order count come in within 20% of the seven-day average?" only means something in the context of today's actual data — data test.

This article is entirely about the first column. If you're looking for the second, dbt's testing docs, Great Expectations, and Soda are the right places to go next — and a mature data platform usually runs both, not one instead of the other.

Part 1: pytest fundamentals

Pytest docs

Installing and writing your first test

uv add --dev pytest

This adds pytest as a development dependency — something your project needs to run its own test suite, but not something anyone installing your package needs. uv writes it into a [dependency-groups] table in pyproject.toml (the standardized format from PEP 735), kept separate from your project's real runtime dependencies:

[dependency-groups]
dev = [
    "pytest>=8.3.4",
]

pytest's core idea: a test is just a function whose name starts with test_, living in a file whose name starts with test_ or ends with _test.py. No boilerplate classes required (though you can use them).

# test_transformations.py
def add_tax(price: float, rate: float = 0.1) -> float:
    return round(price * (1 + rate), 2)

def test_add_tax_applies_default_rate():
    result = add_tax(100)
    assert result == 110.0

Run it with uv run, which executes the command inside the project's managed virtual environment without you ever having to activate one by hand:

uv run pytest
# or, more verbosely:
uv run pytest -v test_transformations.py

pytest uses the plain assert keyword — no self.assertEqual(...) ceremony. When an assertion fails, pytest rewrites it under the hood to show you exactly what was compared, which is a big part of why it's more pleasant than the built-in unittest.

Test discovery and project layout

A typical DE repo looks like:

my_pipeline/
├── src/
│   └── my_pipeline/
│       ├── __init__.py
│       ├── extract.py
│       ├── transform.py
│       └── load.py
├── tests/
│   ├── conftest.py
│   ├── unit/
│   │   ├── test_transform.py
│   │   └── test_extract.py
│   └── integration/
│       └── test_pipeline_end_to_end.py
├── pyproject.toml   # dependencies, dev group, and pytest config all live here
└── uv.lock          # exact resolved versions — commit this to the repo

pytest configuration lives in pyproject.toml too, under [tool.pytest.ini_options] — no separate pytest.ini needed:

[tool.pytest.ini_options]
testpaths = ["tests"]
addopts = "-ra"

testpaths tells pytest (and your editor's test runner) where to look without specifying a path on every invocation; addopts bakes in flags you'd otherwise retype constantly — -ra here prints a one-line summary of every non-passing test at the end of the run.

Separating unit/ and integration/ isn't required, but it lets you run fast tests constantly and slow ones less often:

uv run pytest tests/unit          # fast, run on every save
uv run pytest tests/integration   # slower, run before pushing

Fixtures: pytest's dependency injection

Fixtures are reusable pieces of setup, declared with @pytest.fixture and requested by name as a test function argument. This is the single most important pytest feature for DE work, because pipelines need repeatable inputs — sample dataframes, temp directories, mock connections.

import pytest
import pandas as pd

@pytest.fixture
def raw_orders() -> pd.DataFrame:
    return pd.DataFrame({
        "order_id": [1, 2, 3],
        "customer_id": [10, 10, 11],
        "amount": [25.0, 40.0, None],
    })

def test_drop_nulls_removes_incomplete_rows(raw_orders):
    from my_pipeline.transform import drop_null_amounts
    result = drop_null_amounts(raw_orders)
    assert len(result) == 2
    assert result["amount"].isnull().sum() == 0

Fixtures can depend on other fixtures, and pytest resolves the graph for you. They can also have a scope, controlling how often they're recreated:

@pytest.fixture(scope="function")   # default: fresh per test
@pytest.fixture(scope="module")     # once per test file
@pytest.fixture(scope="session")    # once per whole test run

scope="session" matters a lot for expensive setup — like spinning up a local Spark session (more on this below). The tradeoff is isolation versus speed: a fresh fixture per test can never leak state between tests, while a shared one is faster but puts the burden on you to make sure nothing one test does lingers to affect the next.

conftest.py: sharing fixtures across files

Fixtures defined in tests/conftest.py are automatically available to every test file in that directory and below, no import needed. This is where you put your "standard" sample datasets, temp-directory helpers, and mock clients so every test file can reuse them without duplication.

Parametrize: one test, many cases

Data pipelines are full of edge cases — empty strings, nulls, negative numbers, weird encodings. @pytest.mark.parametrize lets you run the same test logic against a table of inputs and expected outputs instead of copy-pasting near-identical test functions.

import pytest

@pytest.mark.parametrize("raw,expected", [
    ("2024-01-15", "2024-01-15"),
    ("2024/01/15", "2024-01-15"),
    ("15-01-2024", "2024-01-15"),
    ("", None),
    (None, None),
])
def test_normalize_date(raw, expected):
    from my_pipeline.transform import normalize_date
    assert normalize_date(raw) == expected

This is arguably the highest-leverage pytest feature for DE testing: it forces you to explicitly enumerate the messy input variants you actually expect from real-world data, instead of testing only the happy path.

Marks: skip, xfail, and custom categories

import pytest
import sys

@pytest.mark.skipif(sys.platform == "win32", reason="path handling differs on Windows")
def test_partition_path_format():
    ...

@pytest.mark.slow
def test_full_backfill_pipeline():
    ...

Register custom marks like slow or spark in the same [tool.pytest.ini_options] table from the project layout above, so there's one config file for the whole project instead of two competing ones:

# pyproject.toml
[tool.pytest.ini_options]
testpaths = ["tests"]
markers = [
    "slow: long-running integration tests",
    "spark: tests requiring a SparkSession",
]

uv run pytest -m "not slow"     # skip slow tests during local dev
uv run pytest -m spark          # run only Spark tests

(A standalone pytest.ini file still works if you'd rather keep pytest's config out of pyproject.toml — but there's little reason to when everything else about the project, dependencies included, already lives there.)

Mocking

unittest.mock (built into the standard library) lets you replace a real dependency — an API call, a database connection, an S3 client — with a fake that returns canned data. This keeps unit tests fast and independent of network or infrastructure.

from unittest.mock import patch, MagicMock

def test_fetch_exchange_rate_handles_api_response():
    from my_pipeline.extract import fetch_exchange_rate

    fake_response = MagicMock()
    fake_response.json.return_value = {"rate": 1.08}
    fake_response.status_code = 200

    with patch("my_pipeline.extract.requests.get", return_value=fake_response):
        rate = fetch_exchange_rate("USD", "EUR")

    assert rate == 1.08

The key discipline: mock at the boundary of your system, not deep inside your own logic. If you find yourself mocking three layers deep to test a transformation function, that's usually a sign the function is doing too much and should be split into a pure part (testable without mocks) and an I/O part (tested with mocks or integration tests).

Part 2: the AAA pattern

Arrange–Act–Assert is a structural convention, not a pytest feature, but it keeps tests readable as your suite grows:

def test_deduplicate_orders_keeps_latest_record():
    # Arrange
    df = pd.DataFrame({
        "order_id": [1, 1, 2],
        "updated_at": ["2024-01-01", "2024-01-05", "2024-01-01"],
        "status": ["pending", "shipped", "pending"],
    })

    # Act
    result = deduplicate_orders(df)

    # Assert
    assert len(result) == 2
    assert result.loc[result.order_id == 1, "status"].iloc[0] == "shipped"

Every test should have exactly one clear "Act" step and assertions that check one behavior, even if that takes multiple assert lines. If a test's Arrange section is enormous and its Assert section is checking five unrelated things, split it — you'll thank yourself when it fails and you need to know why in five seconds, not five minutes.

Part 3: unit vs. integration tests, for pipelines specifically

The unit/integration distinction maps onto DE work a bit differently than it does onto typical application code:

Unit tests — test a single transformation function in isolation, with an in-memory dataframe you constructed by hand. No file I/O, no database, no cluster. These should run in milliseconds and make up the bulk of your suite.

def test_calculate_discount_caps_at_50_percent():
    result = calculate_discount(loyalty_years=20, base_rate=0.05)
    assert result == 0.5

Integration tests — test that multiple pieces work together against something closer to real infrastructure: a real (but local/test) database, a real file read/write, a local Spark session, a mocked-but-realistic S3 bucket (via moto). These are slower and fewer in number, but they catch the bugs unit tests structurally can't — a SQL query that's syntactically valid but returns the wrong join cardinality, a Parquet schema mismatch between writer and reader.

def test_pipeline_writes_expected_row_count(tmp_path, raw_orders):
    from my_pipeline.pipeline import run_pipeline

    input_path = tmp_path / "orders.csv"
    output_path = tmp_path / "output.parquet"
    raw_orders.to_csv(input_path, index=False)

    run_pipeline(input_path, output_path)

    result = pd.read_parquet(output_path)
    assert len(result) == 2

tmp_path is a built-in pytest fixture that gives you a fresh temporary directory per test, auto-cleaned afterward — extremely useful for testing anything that reads or writes files, without polluting your real filesystem or needing manual teardown.

The shape of a healthy suite follows from how expensive each layer is to run and how much of your logic it can realistically cover:

A useful rule of thumb: if you can't explain in one sentence what real-world bug a test would catch, it's probably testing implementation detail rather than behavior — cut it or rewrite it.

Part 4: testing pandas and Polars pipelines

pandas

The standard library ships purpose-built comparison helpers — use them instead of ==, because dataframe equality has edge cases (dtype mismatches, index alignment, float precision) that == handles inconsistently.

import pandas as pd
from pandas.testing import assert_frame_equal, assert_series_equal

def test_aggregate_revenue_by_region():
    input_df = pd.DataFrame({
        "region": ["west", "west", "east"],
        "revenue": [100, 200, 50],
    })

    result = aggregate_revenue_by_region(input_df)

    expected = pd.DataFrame({
        "region": ["east", "west"],
        "revenue": [50, 300],
    })
    assert_frame_equal(
        result.reset_index(drop=True),
        expected.reset_index(drop=True),
        check_dtype=False,
    )

check_dtype=False is worth knowing about: it's common for a groupby-aggregate to return int64 where your hand-built expected frame has int64 too, but small differences (e.g., float64 vs float32) shouldn't fail a test that's really checking values, not storage format — unless dtype correctness is exactly what you're testing, in which case leave it on.

Polars

Polars ships an equivalent testing module:

import polars as pl
from polars.testing import assert_frame_equal

def test_filter_active_customers():
    df = pl.DataFrame({
        "customer_id": [1, 2, 3],
        "is_active": [True, False, True],
    })

    result = filter_active_customers(df)

    expected = pl.DataFrame({
        "customer_id": [1, 3],
        "is_active": [True, True],
    })
    assert_frame_equal(result, expected)

Because Polars encourages a lazy/expression-based style, it's often cleanest to test the underlying expression logic directly (e.g., a function that returns a pl.Expr) separately from the I/O that triggers .collect(). That keeps the fast unit-testable core small and pure.

Property-based testing (optional but powerful)

For transformation logic with many edge cases, hypothesis can generate hundreds of varied inputs automatically instead of you hand-writing each case:

from hypothesis import given
from hypothesis import strategies as st

@given(st.lists(st.floats(allow_nan=False, allow_infinity=False)))
def test_normalize_never_produces_values_outside_zero_one(values):
    if not values:
        return
    result = min_max_normalize(values)
    assert all(0 <= v <= 1 for v in result)

This is worth reaching for once your parametrize list starts feeling like you're guessing at edge cases rather than enumerating known ones.

Part 5: testing PySpark pipelines

Spark's biggest testing challenge is the startup cost of a SparkSession. Solve it with a session-scoped fixture so it's created once for the whole test run, not once per test:

# tests/conftest.py
import pytest
from pyspark.sql import SparkSession

@pytest.fixture(scope="session")
def spark():
    spark = (
        SparkSession.builder
        .master("local[2]")
        .appName("pytest-spark")
        .config("spark.sql.shuffle.partitions", "2")  # keep local runs fast
        .getOrCreate()
    )
    yield spark
    spark.stop()

One gotcha with sharing a session across every test: anything a test leaves behind — a temp view, a changed config, a cached table — is still there for the next test, since they're the same session. If tests start passing or failing depending on execution order, that's usually the tell. A cheap guard is a small autouse fixture that clears temp views between tests, or dropping to scope="module" for the specific test file where isolation matters more than the extra setup cost.

Every test that needs Spark just requests the spark fixture:

def test_flag_high_value_orders(spark):
    df = spark.createDataFrame(
        [(1, 500.0), (2, 50.0)],
        ["order_id", "amount"],
    )

    result = flag_high_value_orders(df, threshold=100.0)

    rows = {r.order_id: r.is_high_value for r in result.collect()}
    assert rows == {1: True, 2: False}

Two libraries make Spark dataframe assertions much less painful than manual .collect() comparisons:

• chispa — gives you assert_df_equality(result, expected, ignore_row_order=True) with readable diff output, similar in spirit to pandas.testing.assert_frame_equal. It's the most established option and still the one with the nicest failure messages.
• pytest-spark — provides Spark-related fixtures and config out of the box if you don't want to hand-roll the session fixture above.

If you're on Spark 4.0+, it's also worth knowing PySpark now ships a built-in pyspark.testing.assertDataFrameEqual, so you can get row/column-order-insensitive comparisons without a third-party dependency. chispa still has the edge on diff readability, but the native option is a reasonable default if you'd rather not add a dependency.

from chispa import assert_df_equality

def test_join_orders_with_customers(spark):
    orders = spark.createDataFrame([(1, 10)], ["order_id", "customer_id"])
    customers = spark.createDataFrame([(10, "Acme")], ["customer_id", "name"])

    result = join_orders_with_customers(orders, customers)

    expected = spark.createDataFrame([(1, 10, "Acme")], ["order_id", "customer_id", "name"])
    assert_df_equality(result, expected, ignore_row_order=True, ignore_column_order=True)

For Spark specifically, mark these tests (@pytest.mark.spark) and consider keeping them out of the default fast test run — a two-node local session still takes a few seconds to spin up, which adds up across a large suite.

Part 6: mocking external systems

DE pipelines are full of edges that touch the outside world: S3, a warehouse, a REST API, a message queue. You don't want unit tests hitting real infrastructure — it's slow, flaky, and sometimes destructive.

moto mocks AWS services at the boundary so your code calls the real boto3 API but nothing actually leaves your machine:

import boto3
from moto import mock_aws

@mock_aws
def test_upload_writes_expected_key():
    s3 = boto3.client("s3", region_name="us-east-1")
    s3.create_bucket(Bucket="test-bucket")

    from my_pipeline.load import upload_parquet
    upload_parquet(s3, bucket="test-bucket", key="orders/2024-01-01.parquet", data=b"fake-bytes")

    objects = s3.list_objects_v2(Bucket="test-bucket")
    keys = [o["Key"] for o in objects["Contents"]]
    assert "orders/2024-01-01.parquet" in keys

For databases, prefer a real-but-disposable instance over mocking the driver whenever practical — e.g., SQLite in-memory for logic that's DB-agnostic, or a Dockerized Postgres/test schema for integration tests that need to check actual SQL behavior. Mocking a database connection to return canned fetchall() results tests your Python glue code, but it can't catch a broken JOIN or a typo in a column name — only a real query engine can.

The clock is an external dependency too, and it's an easy one to forget. Any pipeline logic that reasons about "today," "yesterday's partition," or "records from the last 24 hours" is implicitly depending on datetime.now() — which means the test's outcome depends on when you happen to run it, unless you pin it down. freezegun (or the newer time-machine, which does the same job faster) fixes this by mocking the clock itself:

from freezegun import freeze_time

def test_partition_path_uses_yesterday():
    from my_pipeline.extract import yesterday_partition_path

    with freeze_time("2024-03-15"):
        assert yesterday_partition_path() == "orders/dt=2024-03-14"

Without this, a test like the one above either hardcodes today's date (and quietly breaks tomorrow) or skips testing the date logic entirely — both worse options than mocking the one dependency that's actually causing the problem.

Part 7: data-quality-specific testing patterns

A few habits particular to data engineering that don't show up in typical backend testing:

Test schemas, not just values. A transformation can return the "right" numbers with the wrong column names or types, and a naive test that only checks a couple of cell values will miss it. Libraries like pandera (for pandas/Polars) let you assert against a schema as part of your test:

import pandera as pa
from pandera import Column, DataFrameSchema

order_schema = DataFrameSchema({
    "order_id": Column(int, unique=True),
    "amount": Column(float, pa.Check.ge(0)),
})

def test_transform_output_matches_schema(raw_orders):
    result = transform_orders(raw_orders)
    order_schema.validate(result)  # raises if schema doesn't match

This is the pandera use case from the code-testing column of the framework earlier in this article: a schema checked against a fixture you built by hand, inside a pytest test, as part of CI. The same order_schema object can just as easily validate a real dataframe pulled from production at pipeline runtime — at that point it's stopped being a code test and become a data test, even though the schema definition didn't change. Worth remembering which hat it's wearing in a given call site.

Build small, deliberately ugly fixtures. Real production data is a bad test fixture — it's huge, it changes, and it obscures which specific case you're testing. A handful of hand-built rows that include a null, a duplicate, an empty string, and a negative number will catch more bugs than a 10,000-row sample of "normal" data ever will.

Separate pure transformation logic from I/O. A function like def transform(df: pd.DataFrame) -> pd.DataFrame is trivially unit-testable. A function like def run(): df = pd.read_csv(...); ...; df.to_sql(...) is not — you're forced into slow integration tests for everything. Structure pipelines as thin I/O wrappers around pure, well-tested transformation functions wherever you can.

Test failure paths, not just success paths. What happens when the input file is empty? When a required column is missing? When two upstream systems disagree on a foreign key? These are the tests that actually save you in production, and they're the ones people skip because writing the happy-path test already "felt done."

Use golden/reference datasets for complex aggregations. For a business-logic-heavy transformation (e.g., a multi-step revenue reconciliation), it's sometimes more maintainable to check a small input CSV against a small expected-output CSV committed to the repo, rather than constructing dataframes inline in every test.

Part 8: wiring it into CI

None of this pays off if it only runs on your laptop. A minimal GitHub Actions setup:

# .github/workflows/test.yml
name: tests
on: [push, pull_request]
jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - name: Install uv
        uses: astral-sh/setup-uv@v7
        with:
          python-version: "3.11"
          enable-cache: true
      - run: uv sync --locked --all-extras --dev
      - run: uv run pytest tests/unit -v
      - run: uv run pytest tests/integration -v -m "not slow"

uv sync --locked installs exactly what's pinned in uv.lock and fails the build if the lockfile is out of date with pyproject.toml — the CI equivalent of "works on my machine" actually meaning something. astral-sh/setup-uv is the official action for installing uv itself; it can also pin the Python version the same way actions/setup-python used to, and enable-cache: true caches uv's package store between runs so later builds skip re-downloading dependencies that haven't changed.

Keeping unit and integration runs as separate steps means a failing integration test doesn't hide a failing unit test in the same log, and you get faster feedback from the unit step first.

Putting it together: a minimal but real test suite

tests/
├── conftest.py              # shared fixtures: sample dfs, spark session, tmp helpers
├── unit/
│   ├── test_transform.py    # pure functions, pandas/polars asserts, parametrize-heavy
│   └── test_validation.py   # schema checks, edge cases
└── integration/
    ├── test_spark_jobs.py   # marked @pytest.mark.spark, uses chispa
    ├── test_s3_io.py        # uses moto
    └── test_end_to_end.py   # runs the full pipeline against tmp_path input/output

The core habits worth taking away

1. Know which kind of testing you're doing. pytest catches bugs in your code, using data you control. It cannot catch a data quality problem that only exists in today's actual data — that's a separate job for dbt test, Great Expectations, or Soda.
2. Test the transformation logic, not the framework. You don't need to test that pandas' groupby works — you need to test that your aggregation logic does the right thing.
3. Keep unit tests fast and dependency-free; push anything touching a real file, database, or cluster into a clearly separated, clearly marked integration suite.
4. Build small, deliberately messy fixtures instead of testing against production-sized samples.
5. Use the domain-specific assertion helpers (assert_frame_equal, assert_df_equality, pandera schemas) instead of hand-rolled comparisons — they exist because naive equality checks on dataframes are full of footguns.
6. Mock at the boundary, not inside your own logic — and remember the clock counts as a boundary too. If mocking feels awkward, it's often telling you to refactor, not to mock harder.

None of this requires a big investment up front. Start by putting AAA-structured unit tests around your messiest transformation function, get comfortable with fixtures and parametrize, and expand outward from there. The payoff compounds fast: every bug a test catches before a stakeholder sees a bad number is time you don't spend doing forensic debugging on a Friday afternoon.