DE Mastery GitHub

Apache Spark · SPK-12

Testing and production Spark

From your world. dbt tests and CI

← All exercises

Write tests that catch the bugs

This exercise reverses the usual task: you write the test. A revenue pipeline computes the total line_amount per company from an order_items-style input. You must write the assertion function that decides whether the result of a pipeline is correct. It must be strong enough to pass a correct pipeline and to catch the buggy ones.

Implement assert_valid(spark, result_df, orders_df) in solution.py:

  • orders_df is the raw input. It has the columns company_id, order_item_id, line_amount (with some duplicate rows, as real feeds have).
  • result_df is the output of some pipeline. The pipeline claims that it is the revenue per company. It has the columns company_id, revenue.
  • The function must return normally when result_df is a correct per-company revenue summary of orders_df. It must raise AssertionError when the result is wrong.

To be trustworthy, your checks must not trust the numbers of the result. Recompute the correct answer independently (an oracle, with Spark or DuckDB over orders_df). Compare it row for row (order-insensitive, float tolerance). Add the invariants that a correct result must satisfy:

  • one row per company
  • exactly the companies of the input
  • no NULL keys
  • non-negative revenue

The grader runs your assert_valid against a known-good pipeline (it must pass). It also runs it against several planted buggy pipelines (your checks must catch each one). For example:

  • one pipeline sums only distinct rows (it drops the revenue of true duplicate lines)
  • one pipeline averages instead of summing
  • one pipeline drops a company
  • one pipeline returns negative/NULL revenue

A suite that misses any of these is too weak. This is exactly the failure mode that this module is about.

Starter code

solution.py
"""SPK-12 / 01-test-the-pipeline — starter. See README.md.

Write assert_valid so it PASSES a correct per-company revenue result and RAISES
AssertionError on a wrong one. This starter is too weak: it only checks the row count,
so it fails to catch most bugs.
"""


def assert_valid(spark, result_df, orders_df):
    # TOO WEAK: only checks that there is at least one row. Replace with a real oracle
    # comparison plus invariants (see README.md).
    assert result_df.count() >= 1

Hints

Hint 1

You are given a pipeline's RESULT plus the raw ORDERS it was computed from. Independently recompute the correct per-company revenue (an oracle) and assert the result equals it — row-level, order-insensitive, float tolerance (pyspark.testing.assertDataFrameEqual, checkRowOrder=False).

Hint 2

Also assert invariants that a correct result must satisfy: one row per company, no NULL company_id, revenue is non-negative, and the companies present match the input's companies. These catch bugs an equality check alone might miss (e.g. a dropped company).

Hint 3

assert_valid(spark, result_df, orders_df) must RAISE (AssertionError) when the result is wrong and return normally when it is correct. Build the oracle with Spark or DuckDB over orders_df; don't trust the result's own numbers.

Run it locally
uv run de start SPK-12/01-test-the-pipeline
# edit workspace/SPK-12/01-test-the-pipeline/
uv run de test