Module 14: Testing & Data Quality Engineering

# test_transformations.py
import pytest
from pyspark.sql import SparkSession, functions as F
from pyspark.testing import assertDataFrameEqual
from transformations import clean_orders

COLS = "order_id INT, customer_id INT, product_id INT, quantity INT, status STRING, order_ts STRING"

@pytest.fixture(scope="session")
def spark():
    return SparkSession.builder.getOrCreate()

def make(spark, rows):
    return spark.createDataFrame(rows, COLS).withColumn("order_ts", F.to_timestamp("order_ts"))

def test_status_is_lowercased_and_trimmed(spark):
    out = clean_orders(make(spark, [(1, 10, 5, 2, "  COMPLETED ", "2025-03-01 10:00:00")]))
    assert out.first()["status"] == "completed"

def test_missing_customer_is_invalid(spark):
    out = clean_orders(make(spark, [(1, None, 5, 2, "shipped", "2025-03-01 10:00:00")]))
    assert out.first()["is_valid"] is False

def test_zero_quantity_is_invalid(spark):
    out = clean_orders(make(spark, [(1, 10, 5, 0, "shipped", "2025-03-01 10:00:00")]))
    assert out.first()["is_valid"] is False

def test_full_output(spark):
    out = clean_orders(make(spark, [(1, 10, 5, 2, "Shipped", "2025-03-01 10:00:00")]))
    expected = spark.createDataFrame(
        [(1, 10, 5, 2, "shipped", True)],
        "order_id INT, customer_id INT, product_id INT, quantity INT, status STRING, is_valid BOOLEAN",
    )
    cols = ["order_id", "customer_id", "product_id", "quantity", "status", "is_valid"]
    assertDataFrameEqual(out.select(*cols), expected)

%pip install pytest

import sys, pytest
sys.dont_write_bytecode = True
retcode = pytest.main([".", "-v", "-p", "no:cacheprovider"])
assert retcode == 0, "Tests failed: see the log"

catalog = spark.sql("SELECT current_catalog()").first()[0]
test_schema = f"{catalog}.dbx_test"
spark.sql(f"CREATE SCHEMA IF NOT EXISTS {test_schema}")

# 1. arrange: a known input table
(spark.table("orders").limit(100)
      .write.mode("overwrite").saveAsTable(f"{test_schema}.orders_in"))

# 2. act: run the real code against it
from transformations import clean_orders
clean_orders(spark.table(f"{test_schema}.orders_in")) \
    .write.mode("overwrite").saveAsTable(f"{test_schema}.orders_out")

# 3. assert
out = spark.table(f"{test_schema}.orders_out")
assert out.count() == 100, "rows were lost or added"
assert out.filter("status != lower(trim(status))").count() == 0
assert out.select("order_id").distinct().count() == 100, "order_id must stay unique"

from pyspark.sql import functions as F

def fake_orders(spark, n=1000):
    """Same rows every run: values derive from the id with a hash."""
    return spark.range(1, n + 1).select(
        F.col("id").cast("int").alias("order_id"),
        (F.pmod(F.hash("id", F.lit(1)), F.lit(50)) + 1).cast("int").alias("customer_id"),
        (F.pmod(F.hash("id", F.lit(2)), F.lit(30)) + 1).cast("int").alias("product_id"),
        (F.pmod(F.hash("id", F.lit(3)), F.lit(5)) + 1).cast("int").alias("quantity"),
        F.lit(" Completed ").alias("status"),
        F.timestamp_seconds(F.lit(1740787200) + F.col("id") * 60).alias("order_ts"),
    )

from pyspark.testing import assertSchemaEqual

expected = spark.createDataFrame([], "order_id INT, customer_id INT, product_id INT, quantity INT, status STRING, order_ts TIMESTAMP").schema
actual = spark.table("orders").select("order_id", "customer_id", "product_id", "quantity", "status", "order_ts").schema
assertSchemaEqual(actual, expected)

contract = {
    "table": "orders",
    "columns": {"order_id": "int", "customer_id": "int", "quantity": "int"},
    "not_null": ["order_id"],
    "unique": ["order_id"],
}

df = spark.table(contract["table"])
types = dict(df.dtypes)
for col, t in contract["columns"].items():
    assert types.get(col) == t, f"{col}: expected {t}, got {types.get(col)}"
for col in contract["not_null"]:
    assert df.filter(F.col(col).isNull()).count() == 0, f"{col} has nulls"
for col in contract["unique"]:
    assert df.count() == df.select(col).distinct().count(), f"{col} is not unique"

%sql
ALTER TABLE orders_delta ADD CONSTRAINT qty_positive CHECK (quantity > 0 OR quantity IS NULL)

rules = {
    "order_id not null":    "order_id IS NOT NULL",
    "quantity positive":    "quantity IS NULL OR quantity > 0",
    "status known":         "status IN ('completed', 'shipped', 'cancelled', 'returned')",
    "customer exists":      None,       # needs a join: see below
}

def check(df, rules):
    out = []
    for name, cond in rules.items():
        if cond is None:
            continue
        failed = df.filter(f"NOT ({cond})").count()
        out.append((name, df.count(), failed))
    return spark.createDataFrame(out, "rule STRING, rows LONG, failed LONG")

display(check(spark.table("orders"), rules))

# referential check: every customer_id exists in customers
orphans = (spark.table("orders").join(spark.table("customers"), "customer_id", "left_anti")
                .filter("customer_id IS NOT NULL").count())
assert orphans == 0, f"{orphans} orders point to missing customers"

bronze = spark.table("orders_raw")
silver = spark.table("orders")

checks = {
    "no rows lost after dedupe": silver.count() == bronze.select("order_id").distinct().count(),
    "keys match": bronze.select("order_id").distinct().exceptAll(silver.select("order_id")).count() == 0,
    "quantity sums match": (
        silver.agg(F.sum("quantity")).first()[0]
        == bronze.dropDuplicates(["order_id"]).agg(F.sum("quantity")).first()[0]
    ),
}
for name, ok in checks.items():
    print("PASS" if ok else "FAIL", name)
assert all(checks.values())

# 1. one time: save today's output as the baseline
(spark.table("orders").groupBy("status").agg(F.count("*").alias("n"), F.sum("quantity").alias("qty"))
      .write.mode("overwrite").saveAsTable("baseline_orders_by_status"))

# 2. after a code change: recompute and compare
from pyspark.testing import assertDataFrameEqual
current = spark.table("orders").groupBy("status").agg(F.count("*").alias("n"), F.sum("quantity").alias("qty"))
assertDataFrameEqual(current, spark.table("baseline_orders_by_status"))

%sql
SELECT * FROM orders_delta VERSION AS OF 0
EXCEPT
SELECT * FROM orders_delta

def test_scd2_update(spark):
    # arrange
    spark.createDataFrame([(1, "A", "NL", "2024-01-01", None, True)],
        "customer_id INT, name STRING, country STRING, valid_from STRING, valid_to STRING, is_current BOOLEAN"
    ).write.mode("overwrite").saveAsTable("dbx_test.dim_t")
    # act: apply your merge function with the change (customer 1 -> 'XX')
    # apply_scd2(spark, "dbx_test.dim_t", changes_df)
    # assert
    t = spark.table("dbx_test.dim_t")
    assert t.count() == 2
    assert t.filter("is_current").first()["country"] == "XX"
    assert t.filter("customer_id = 1 AND is_current").count() == 1      # exactly one current row

# arrange: empty test landing folder + fresh checkpoint
catalog = spark.sql("SELECT current_catalog()").first()[0]
test_base = f"/Volumes/{catalog}/dbx_course/raw/stream_test"
dbutils.fs.rm(test_base, recurse=True)

(spark.range(0, 100).select("id", F.lit("buy").alias("event_type"))
      .coalesce(1).write.json(f"{test_base}/in/batch_1"))

def run():
    q = (spark.readStream.schema("id LONG, event_type STRING")
           .option("recursiveFileLookup", "true").option("pathGlobFilter", "*.json")
           .json(f"{test_base}/in")
           .writeStream.option("checkpointLocation", f"{test_base}/cp")
           .trigger(availableNow=True).toTable("dbx_test.events_out"))
    q.awaitTermination()
    return spark.table("dbx_test.events_out").count()

assert run() == 100          # first run reads everything
assert run() == 100          # second run: no new files, no duplicates

import time

def run_job_step():
    return (spark.table("orders").join(spark.table("customers"), "customer_id")
                 .groupBy("country").agg(F.count("*").alias("n")).collect())

start = time.time()
run_job_step()
elapsed = time.time() - start

print(f"{elapsed:.1f}s")
assert elapsed < 60, f"step took {elapsed:.0f}s, budget is 60s"