DataLane
← All cheat sheets

pytest for Data Pipelines cheat sheet

Fixtures, parametrization, dataframe assertions, fakes for warehouses and object stores, and the CI settings that keep suites fast.

ProgrammingAdvanced7 sections

Fixtures and scope

@pytest.fixture(scope="session") def spark(): s = SparkSession.builder.master("local[2]").getOrCreate() yield s s.stop()
A Spark or warehouse session costs seconds to build, so pay for it once per session. Anything mutable it holds still has to be reset per test.
@pytest.fixture def warehouse(tmp_path): con = duckdb.connect(tmp_path / "test.db") yield con con.close()
tmp_path is a fresh directory per test and pytest keeps the last three runs on disk, which is what lets you open the artifact after a CI failure.
conftest.py
Fixtures here are visible to every test in the directory and below with no import. Deep conftest trees make fixture origins untraceable, so keep it shallow.
@pytest.fixture(autouse=True)
Applies to every test in scope without being requested. Correct for resetting global state or clearing a cache, wrong for anything expensive.
request.getfixturevalue("warehouse")
Resolves a fixture by name at run time. This is the escape hatch for choosing which backend a test runs against based on a parameter.

Parametrization

@pytest.mark.parametrize( "raw,expected", [("2026-01-31", date(2026, 1, 31)), ("", None)], ids=["iso", "empty"], )
Explicit ids make the failure line readable and give CI a stable node id to rerun. Without them pytest stringifies the arguments into noise.
@pytest.mark.parametrize("engine", ["duckdb", "postgres"], indirect=True)
indirect routes the value into a fixture of the same name, which is how one test body runs unchanged against several real backends.
pytest.param(bad_row, marks=pytest.mark.xfail(strict=True))
strict xfail fails the suite once the case starts passing, so a fixed bug cannot sit forever marked as broken.
@pytest.fixture(params=["parquet", "csv"]) def fmt(request): return request.param
Parametrizing a fixture multiplies every test that uses it. Powerful, and the usual reason a suite suddenly takes four times longer.
pytest --collect-only -q
Prints the expanded matrix without running it. Check this after adding parametrization, since case counts grow multiplicatively rather than additively.

Asserting on data

pd.testing.assert_frame_equal(actual, expected, check_dtype=False, check_like=True)
check_like ignores row and column order, which is what set-semantics SQL output needs. Leave check_dtype on unless int32 versus int64 truly does not matter.
polars.testing.assert_frame_equal(actual, expected, check_row_order=False)
The polars equivalent. Row order after a parallel group by is nondeterministic, so asserting on it produces a test that fails one run in twenty.
assert dict(df.schema) == expected_schema
Assert schema and values in separate tests so an added column and a changed number produce different failures, rather than one unreadable diff.
assert actual == pytest.approx(expected, rel=1e-6)
Float aggregates differ in the last bits depending on partition order. Exact equality on a sum of floats is a guaranteed intermittent failure.
assert df.filter(pl.col("customer_id").is_null()).height == 0, "null keys leaked"
Assert the invariant rather than a snapshot of the output. Invariant tests survive a fixture refresh, golden files do not.

Fakes for warehouses and object stores

monkeypatch.setattr(loader, "read_from_snowflake", lambda *_: fixture_df)
monkeypatch reverts at teardown while a bare setattr leaks into later tests. Patch the name where it is used, not where it is defined.
@mock_aws def test_upload(s3_client): s3_client.create_bucket(Bucket="lake")
moto emulates S3 in process with no network. Version 5 collapsed the per-service decorators into the single mock_aws entry point.
PostgresContainer("postgres:17")
testcontainers gives you a real engine, so constraint violations and type coercion surface in tests. Budget the startup and reuse it across the session.
con.register("orders", df)
DuckDB exposes an in-memory frame as a table, which lets you unit test the exact SQL text of a transformation with no warehouse connection.
responses.add(responses.GET, url, json=payload, status=429)
Stub the HTTP layer rather than your own client wrapper, so retry, backoff, and pagination logic stays inside the tested surface.

Determinism

monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1")
Pin every credential-shaped variable in a fixture. A test that passes only on your laptop is almost always reading an environment variable you forgot.
with freeze_time("2026-08-30T00:00:00Z"): ...
Pins now() so window and partition logic becomes testable. Injecting a clock parameter is better where you control the code, since freezing is global.
monkeypatch.delenv("SNOWFLAKE_PASSWORD", raising=False)
Proves the code fails loudly on a missing secret instead of connecting as some default identity. raising=False keeps it green when the variable was never set.
rng = np.random.default_rng(0)
Seed a generator instance rather than the global random state, so parallel workers do not share or reorder one sequence.
pytest -p no:randomly
pytest-randomly shuffles order to expose hidden dependencies between tests. Disable it only to reproduce one specific failing order.

Selection, markers, and CI

[pytest] markers = integration: hits a live warehouse addopts = -q --strict-markers
strict-markers turns a misspelled marker into an error instead of a test that quietly escapes every filter. Register each marker you use.
pytest -m "not integration"
The default fast path for pull requests. Run the integration set on merge and on a schedule, not on every push to a branch.
pytest --lf --ff
Reruns last failures first. On a suite with a 30-second session fixture this is the difference between a five-second and a five-minute loop.
pytest -n auto --dist loadscope
pytest-xdist parallelism. loadscope keeps a module's tests on one worker, so a session-scoped database fixture is built once per worker instead of per test.
pytest --cov=pipelines --cov-branch --cov-fail-under=80
Branch coverage catches the untested else in every null-handling path. A line-only percentage flatters transformation code badly.

Flakiness and speed

@pytest.mark.timeout(60)
A hung connection otherwise burns the entire CI budget before anyone notices. Set a global timeout in the ini file and override it per test.
@pytest.mark.flaky(reruns=2)
Reruns hide the cause. Defensible as a temporary shield on a network-bound integration test, never on a pure transformation unit test.
@given(st.lists(st.integers(min_value=0), min_size=1)) def test_output_is_non_negative(xs): assert transform(xs) >= 0
Hypothesis explores the input space and shrinks a failure to its minimal case. It finds the empty-partition and single-row bugs fixtures never cover.
pytest --durations=10
Prints the slowest tests. It is nearly always three unmarked integration tests rather than a suite that is uniformly slow.
pytest.importorskip("pyspark")
Skips cleanly when a heavy optional dependency is absent, so the same suite runs on a laptop and in the full CI image without branching.

From DataLane — tutorials at/blog, practice SQL live in theplayground.

↑↓ navigate openesc close