PART 6 · PYTHON FOR DATA ENGINEERING · CHAPTER 34

Testing with pytest and data quality checks

FreeLesson 71 of 78

Unit tests check your code on known inputs. Data quality checks check the data on every run. A pipeline needs both.

TRANSFORM.PYPYTHON 3.12

STATUS = {"delivered": "delivered", "dlvd": "delivered",
          "cancelled": "cancelled",
          "canceled": "cancelled", "returned": "returned"}

def normalize_status(raw):
    key = (raw or "").strip().lower()
    return STATUS.get(key, "unknown")

def to_amount(text):
    value = float(str(text).replace(",", ""))
    if value <= 0:
        raise ValueError(f"amount must be > 0: {value}")
    return round(value, 2)

TEST_TRANSFORM.PYPYTEST

import pytest
from transform import normalize_status, to_amount

@pytest.mark.parametrize("raw, expected", [
    (" DLVD ", "delivered"),
    ("Canceled", "cancelled"),
    (None, "unknown"),
    ("shipped", "unknown"),
])
def test_normalize_status(raw, expected):
    assert normalize_status(raw) == expected

def test_amount_parses_commas():
    assert to_amount("1,198.50") == 1198.5

def test_amount_rejects_negative():
    with pytest.raises(ValueError, match="> 0"):
        to_amount("-50")

Output

test_normalize_status[ DLVD -delivered] PASSED
test_normalize_status[Canceled-cancelled] PASSED
test_normalize_status[None-unknown] PASSED
test_normalize_status[shipped-unknown] PASSED
test_amount_parses_commas PASSED
test_amount_rejects_negative PASSED
6 passed

DATA QUALITY CHECKS THAT RUN EVERY LOADPYTHON 3.12

import json
import pandas as pd

def run_checks(orders, customers):
    known = set(customers.customer_id)
    checks = {
        "order_id unique": orders.order_id.is_unique,
        "no null customer_id": orders.customer_id.notna().all(),
        "status allowed": orders.status.isin(
            ["delivered", "cancelled", "returned"]).all(),
        "amount > 0": (pd.to_numeric(orders.amount,
                       errors="coerce") > 0).all(),
        "customer exists": orders.customer_id.isin(known).all(),
    }
    return [name for name, ok in checks.items() if not ok]

customers = pd.json_normalize(json.load(open("customers.json")))
for path in ["orders.csv", "orders_dirty.csv"]:
    df = pd.read_csv(path)
    failed = run_checks(df, customers)
    print(path, "->", "PASS" if not failed else f"FAIL {failed}")

Output

orders.csv -> PASS
orders_dirty.csv -> FAIL ['no null customer_id', 'status allowed', 'amount > 0', 'customer exists']

Unit test	Integration test	Data quality check
Checks One function on fixed inputs	Steps together on a small real file and a test database	The real data, every run
Runs On every commit, in CI	Before a deploy	Inside the pipeline, after transform
Logic bugs: wrong mapping, bad Catches rounding	Wiring bugs: wrong path, schema mismatch, bad SQL	Source changes: new status, nulls, duplicates, late files
Example here test_normalize_status	run main() on `orders.csv` into a temp sqlite file	run_checks() on orders_dirty.csv

Practice before you move on

Add a parametrized test case for an amount with spaces around it.
Add a data quality check that fails when an order_id appears twice.