Generating Unit and Data-Validation Tests for Pipelines

Tier: Foundation · Workflow area: Testing & data validation

The skill is prompting for behavior-anchored tests and property-based thinking, then auditing that the tests can actually fail.

When to use

• •
  A transform/feature module is heading to production (pairs with pb-09).
• •
  A pipeline broke silently on a schema or distribution change and you are done being surprised.
• •
  You inherit untested pipeline code and need a safety net before touching it.

Skip exhaustive testing for throwaway analysis code; test what will run unattended.

Prerequisites

• •
  pytest; a data-validation framework — pandera (lightweight, code-first) or Great Expectations (heavier, docs/UI) — chosen per team taste.
• •
  The code under test importable as functions (notebook code: extract first, see pb-09).
• •
  Real knowledge of the data's contract: which invariants are guaranteed vs merely usual. The assistant cannot know that amount can legitimately be negative for refunds — you must say so.

The workflow

1. 1.

   Generate unit tests from stated behavior, not from the code alone:

   Copy

   Write pytest tests for build_features(orders, snapshots) below. Behavior
   contract (test THIS, not the implementation):
   - Output has exactly one row per (customer_id, snapshot_date) input pair.
   - Events with event_ts >= snapshot_date never influence any feature.
   - Customers with zero history yield NaN features, not 0, and are not dropped.
   - Refund rows (amount < 0) reduce spend_30d.
   - Column dtypes match this schema: [paste].
   Use small hand-built fixture DataFrames (5-10 rows) with hand-computed
   expected values in comments. No mocking pandas. Include the temporal
   boundary case: one event exactly at snapshot_date.
   

   Supplying the contract is the difference between tests that catch bugs and tests that enshrine them.

2. 2.

   Audit the tests for falsifiability. For each generated test ask: what bug would make this fail? Then prove it — mutate the code (flip < to <=, drop the refund handling) and confirm the suite goes red. A test suite you have never seen fail is a decoration. Assistants help mechanically: Flag any test that would pass even if the function returned its input unchanged.

3. 3.

   Add property-based tests for the invariants that scale. Beyond fixtures, ask for Hypothesis tests on genuine properties: row-count preservation through non-filtering transforms, idempotence (f(f(x)) == f(x) for cleaners), permutation invariance (input row order must not change aggregates), no-NaN-introduction guarantees. Review the generated strategies — overly narrow strategies (st.integers(1, 10)) test nothing.

4. 4.

   Generate the data-validation suite from profile + contract. Feed pb-01's profiling output and let the assistant draft a pandera schema:

   Copy

   From this profiling summary [paste] and these business rules [amount in
   paise, >0 except refunds; status in {placed,shipped,cancelled}; order_ts
   within last 3 years; customer_id never null; <2% null mrr], generate a
   pandera DataFrameSchema with Checks. Distinguish HARD failures (schema,
   nulls in keys) from WARN-level drift checks (null-rate ranges, category
   sets, value quantiles within historical p1-p99 bands).
   

   The hard/warn split is essential: a new country code should page nobody at 3am; a null primary key should stop the pipeline.

5. 5.

   Add distribution-shift sentinels. For model-feeding pipelines: PSI or KS distance of key features vs a stored reference window, positive-class rate bounds, and volume bounds (row count within ±3σ of the day-of-week norm). Assistants generate competent PSI code; verify the binning against the reference sample yourself.

6. 6.

   Wire placement and failure behavior. Validation runs at ingestion (raw contract), post-transform (feature contract), and pre-scoring (model input schema). Decide per check: block, quarantine, or log-and-continue — and generate CI hooks plus a runtime entry point so the same schema guards tests and production runs.

Tool variants

• pandera (code-first): schemas as Python, versioned with the code, minimal ceremony — the default for sklearn/pandas teams.
• Great Expectations: richer ecosystem — data docs, checkpoints, profiling-assisted suite creation; better when non-engineers consume validation results. The same contract prompt generates expectation suites.
• Chat vs agentic: chat suffices for single modules; an agentic tool (Claude Code) can run the mutate-and-verify falsifiability loop (step 2) itself — have it show you each mutation and the resulting red test.

Verification

• •
  Mutation check passed: every test observed failing against at least one seeded bug; no test passes with the identity implementation.
• •
  The temporal boundary fixture (event at exactly snapshot_date) exists and encodes the documented < vs <= decision.
• •
  Validation suite runs green on the last 30 days of real data — and correctly fires on a deliberately corrupted sample (nulled keys, shifted distribution).
• •
  Hard vs warn severity assigned to every check, with an owner for warns.
• •
  Tests run in CI; validation runs in the production DAG, not just in CI.

Pitfalls

• Tautological tests.result = f(x); assert result.equals(f(x)) and mock-everything tests are the default failure mode of test generation; the falsifiability audit exists precisely for this.
• Asserting the current bug. Generating tests purely from existing code freezes its defects as "expected behavior." Contract first, code second.
• Over-tight validation. Bounds set at exactly the observed min/max page you weekly for normal drift; use historical quantile bands and the warn tier.
• Under-tight validation.col >= 0 on a column that "should" be positive misses the day it becomes 99% zeros; pair range checks with distribution sentinels.
• Fixture sprawl. Fifty near-identical generated fixtures nobody understands; demand minimal fixtures with hand-computed expectations in comments.

Restricted-environment variant

• •
  Test generation needs contracts and schemas, not data — safe with hosted assistants as written; fixtures are synthetic by construction.
• •
  Validation thresholds derived from real profiles: compute the profile locally (pb-01 restricted variant) and paste only the aggregate bands.
• •
  Air-gapped: local code models generate serviceable pytest/pandera code from tight contracts; the falsifiability audit becomes even more important — run it without exception.

Success criteria

• •
  A transform module with fixture-based unit tests, ≥2 property-based tests, and a two-tier validation schema — written in under an hour.
• •
  Every test demonstrated falsifiable via mutation.
• •
  The next upstream schema change or distribution shift is caught by the suite before stakeholders notice — the actual ROI event.
• •
  Team members copy your contract-prompt pattern for their own modules.