Title here
Summary here
Lesson 1 - Unit-Testing Tasks
Welcome to Unit-Testing Tasks
Every DAG so far in this course has been tested the same way: write it, save it to .airflow-local/dags/, wait for the scheduler to notice it, trigger a real run, and poll until it finishes. That’s the right way to test that a DAG works — dependencies, trigger rules, retries, all of it. It is a slow, heavy way to test whether one function’s logic is correct, and CityFlow’s validate_trips task has real logic worth testing on its own: five checks that decide, for every single trip, whether it’s clean or quarantined.
By the end of this lesson, you will be able to:
- Pull a task’s decision logic out into a plain, importable function with no Airflow dependency at all
- Write real
assert-based unit tests covering a valid case, multiple invalid cases, and a genuine edge case - Explain, with a real measured number, why testing logic in isolation is faster than testing it by triggering a DAG
Pulling the Logic Out, Not Copying It
cityflow_etl_dag.py’s validate_trips task (Module 5 Lesson 3) does two jobs at once: it decides whether each row is clean, and it reads/writes the files that make that decision visible. Only the first job is worth unit-testing — the second is just I/O. cityflow_validation_logic.py keeps only the first job, as a plain function:
"""Pure, importable CityFlow trip-validation logic (Module 6 Lesson 1).
This module has no Airflow import at all -- it is the same five real checks
cityflow_etl_dag.py's validate_trips task already runs (null required fields,
out-of-range trip_distance / fare_amount / passenger_count, and a
pickup-after-dropoff ordering check), pulled out into one plain function so it
can be unit-tested in isolation: no scheduler, no worker, no live Airflow
install required at all. cityflow_etl_instrumented_dag.py's validate_trips
task (Lesson 5) imports and calls this exact function -- the real DAG task
becomes a thin wrapper (read the CSV, call this, write the report and
quarantine files) around logic that was already proven correct here, not a
second copy of the same five checks."""
import pandas as pd
REQUIRED_FIELDS = ["passenger_count", "trip_distance", "fare_amount", "PULocationID", "DOLocationID"]
def validate_trips_df(trips: pd.DataFrame) -> tuple[dict, "pd.Series"]:
"""Runs the five real CityFlow data-quality checks against `trips` and
returns (report, any_fail_mask): a JSON-safe report dict shaped exactly
like validate_trips's own validation_report.json, plus the boolean
Series a caller uses to split clean vs. quarantine rows."""
null_mask = trips[REQUIRED_FIELDS].isna().any(axis=1)
dist_mask = (trips["trip_distance"] < 0) | (trips["trip_distance"] > 100)
fare_mask = (trips["fare_amount"] < 0) | (trips["fare_amount"] > 500)
pax_mask = (trips["passenger_count"] < 0) | (trips["passenger_count"] > 6)
time_mask = trips["tpep_pickup_datetime"] > trips["tpep_dropoff_datetime"]
checks = {
"null_required_fields": int(null_mask.sum()),
"trip_distance_out_of_range": int(dist_mask.sum()),
"fare_amount_out_of_range": int(fare_mask.sum()),
"passenger_count_out_of_range": int(pax_mask.sum()),
"pickup_after_dropoff": int(time_mask.sum()),
}
any_fail = null_mask | dist_mask | fare_mask | pax_mask | time_mask
rows_failed = int(any_fail.sum())
rows_passed = int((~any_fail).sum())
pass_rate = rows_passed / len(trips) if len(trips) else 0.0
report = {
"total_rows": len(trips), "checks": checks,
"rows_failed_any_check": rows_failed, "rows_passed": rows_passed,
"pass_rate": round(pass_rate, 4),
}
return report, any_failThis file is saved to .airflow-local/dags/cityflow_validation_logic.py — inside the DAGs folder, but with no @dag decorator, so the scheduler never treats it as a DAG. It sits there as a plain, importable module any real DAG file can pull from, exactly the way Lesson 5’s guided project will.
Why this split matters
validate_trips_df takes a DataFrame and returns a dict — nothing about it knows what Airflow is. That’s the whole point: a function that doesn’t depend on a scheduler doesn’t need one running to be tested.
Seven Real Tests, No Airflow Required
import sys
import time
import pandas as pd
sys.path.insert(0, ".airflow-local/dags")
import cityflow_validation_logic as vlogic
REQUIRED_COLS = ["tpep_pickup_datetime", "tpep_dropoff_datetime", "passenger_count", "trip_distance", "fare_amount", "PULocationID", "DOLocationID"]
def row(**overrides):
base = dict(
tpep_pickup_datetime=pd.Timestamp("2024-01-01 08:00:00"),
tpep_dropoff_datetime=pd.Timestamp("2024-01-01 08:15:00"),
passenger_count=1, trip_distance=3.2, fare_amount=14.5,
PULocationID=142, DOLocationID=236,
)
base.update(overrides)
return base
def df_of(*rows):
return pd.DataFrame(list(rows))[REQUIRED_COLS]
t_start = time.perf_counter()
n_tests = 0
# Test 1: a genuinely valid row passes every check.
report, any_fail = vlogic.validate_trips_df(df_of(row()))
assert report["rows_passed"] == 1 and report["rows_failed_any_check"] == 0
assert not any_fail.iloc[0]
assert all(v == 0 for v in report["checks"].values())
n_tests += 1
print("PASS: a genuinely valid row passes all five checks")
# Test 2: a genuinely invalid row (out-of-range fare_amount) is caught.
report, any_fail = vlogic.validate_trips_df(df_of(row(fare_amount=-12.0)))
assert report["rows_failed_any_check"] == 1
assert report["checks"]["fare_amount_out_of_range"] == 1
assert any_fail.iloc[0]
n_tests += 1
print("PASS: a negative fare_amount is caught by fare_amount_out_of_range")
# Test 3: a genuinely invalid row (a null required field) is caught.
report, any_fail = vlogic.validate_trips_df(df_of(row(passenger_count=None)))
assert report["checks"]["null_required_fields"] == 1
assert any_fail.iloc[0]
n_tests += 1
print("PASS: a null passenger_count is caught by null_required_fields")
# Test 4: a genuinely invalid row (pickup after dropoff) is caught.
report, any_fail = vlogic.validate_trips_df(df_of(row(
tpep_pickup_datetime=pd.Timestamp("2024-01-01 08:30:00"),
tpep_dropoff_datetime=pd.Timestamp("2024-01-01 08:15:00"),
)))
assert report["checks"]["pickup_after_dropoff"] == 1
n_tests += 1
print("PASS: a pickup timestamp after its own dropoff is caught by pickup_after_dropoff")
# Edge case A: the range check is a strict '>', so exactly 100.0 miles passes
# and 100.01 miles fails -- the boundary itself, not just an obviously-bad value.
report, any_fail = vlogic.validate_trips_df(df_of(row(trip_distance=100.0)))
assert report["checks"]["trip_distance_out_of_range"] == 0, "100.0 miles is exactly the boundary and must still pass"
report, any_fail = vlogic.validate_trips_df(df_of(row(trip_distance=100.01)))
assert report["checks"]["trip_distance_out_of_range"] == 1, "100.01 miles is one hundredth of a mile over and must fail"
n_tests += 1
print("PASS: the trip_distance boundary itself (100.0 passes, 100.01 fails) behaves correctly")
# Edge case B: a row that fails TWO checks at once is counted once in
# rows_failed_any_check, but each individual check still reports its own hit --
# the "any" semantics that make the checks dict and the total agree on purpose.
report, any_fail = vlogic.validate_trips_df(df_of(row(fare_amount=-5.0, passenger_count=9)))
assert report["rows_failed_any_check"] == 1, "one bad row must count once, not twice, in the combined total"
assert report["checks"]["fare_amount_out_of_range"] == 1
assert report["checks"]["passenger_count_out_of_range"] == 1
n_tests += 1
print("PASS: a row failing two checks at once counts once in rows_failed_any_check, twice across the two checks")
# Edge case C: an exact regression check against real, already-published numbers.
# This is the same real 2024-01-01 CityFlow slice Module 5 Lesson 3 ingested and
# validated with the original inline logic -- if this pure, extracted function is
# faithful to that logic, it must reproduce those exact real published numbers.
raw = pd.read_parquet(
".airflow-local/dags/data/yellow_tripdata_2024-01.parquet",
columns=["tpep_pickup_datetime", "tpep_dropoff_datetime", "passenger_count", "trip_distance", "PULocationID", "DOLocationID", "fare_amount"],
)
raw["tpep_pickup_datetime"] = pd.to_datetime(raw["tpep_pickup_datetime"])
raw["tpep_dropoff_datetime"] = pd.to_datetime(raw["tpep_dropoff_datetime"])
day = raw[(raw["tpep_pickup_datetime"] >= pd.Timestamp("2024-01-01")) & (raw["tpep_pickup_datetime"] < pd.Timestamp("2024-01-02"))]
report, any_fail = vlogic.validate_trips_df(day)
print(f"regression check against real 2024-01-01 CityFlow data: {report}")
assert report["total_rows"] == 81013, report["total_rows"]
assert report["rows_passed"] == 69290, report["rows_passed"]
assert report["pass_rate"] == 0.8553, report["pass_rate"]
n_tests += 1
print("PASS: the extracted function reproduces Module 5 Lesson 3's real, already-published numbers exactly")
t_end = time.perf_counter()
elapsed_ms = (t_end - t_start) * 1000
print(f"\n{n_tests} real test cases, {elapsed_ms:.1f} ms total wall-clock time (including reading the real ~50MB parquet file for the regression check)")PASS: a genuinely valid row passes all five checks
PASS: a negative fare_amount is caught by fare_amount_out_of_range
PASS: a null passenger_count is caught by null_required_fields
PASS: a pickup timestamp after its own dropoff is caught by pickup_after_dropoff
PASS: the trip_distance boundary itself (100.0 passes, 100.01 fails) behaves correctly
PASS: a row failing two checks at once counts once in rows_failed_any_check, twice across the two checks
regression check against real 2024-01-01 CityFlow data: {'total_rows': 81013, 'checks': {'null_required_fields': 10518, 'trip_distance_out_of_range': 1, 'fare_amount_out_of_range': 1307, 'passenger_count_out_of_range': 4, 'pickup_after_dropoff': 2}, 'rows_failed_any_check': 11723, 'rows_passed': 69290, 'pass_rate': 0.8553}
PASS: the extracted function reproduces Module 5 Lesson 3's real, already-published numbers exactly
7 real test cases, 200.8 ms total wall-clock time (including reading the real ~50MB parquet file for the regression check)Six of these seven tests run against tiny, hand-built one-row DataFrames — that’s on purpose, and it’s most of why this is fast. The seventh deliberately does something heavier (reading the real ~50MB parquet file) specifically to prove the extracted function isn’t just plausible, it’s identical: 81013 total rows, 69290 passed, 0.8553 pass rate — the exact numbers Module 5 Lesson 3 already published from the original inline logic. Even with that real file read included, the whole suite finished in 200.8 ms.
The Real Cost of Testing Through a DAG Instead
The comparison this lesson is actually making needs a real number on the other side, not an assumption. cityflow_hello (Module 2’s three-task DAG) does almost nothing — three tasks, roughly two seconds of real work — which makes it a fair, minimal stand-in for “the cheapest possible real DAG trigger”:
import json
import subprocess
import time
dag_id = "cityflow_hello"
subprocess.run(["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "unpause", dag_id],
capture_output=True, text=True, check=True)
run_id = f"m6l1_timing_{int(time.time())}"
t0 = time.perf_counter()
trigger = subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "trigger", dag_id, "--run-id", run_id],
capture_output=True, text=True,
)
if trigger.returncode != 0:
subprocess.run(["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "reserialize"],
capture_output=True, text=True, check=True)
subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "trigger", dag_id, "--run-id", run_id],
capture_output=True, text=True, check=True,
)
print(f"triggered {dag_id!r} as {run_id!r}")
deadline = time.time() + 90
state = "queued"
while state not in ("success", "failed") and time.time() < deadline:
time.sleep(1)
r = subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "list-runs", "-d", dag_id, "--output", "json"],
capture_output=True, text=True, check=True,
)
runs = json.loads(r.stdout)
match = [x for x in runs if x["run_id"] == run_id]
state = match[0]["state"] if match else "queued"
t1 = time.perf_counter()
print(f"run {run_id!r} reached state: {state}")
assert state == "success"
print(f"real wall-clock time from trigger command to detected success (1s poll interval): {t1 - t0:.2f} s")triggered 'cityflow_hello' as 'm6l1_timing_1784925570'
run 'm6l1_timing_1784925570' reached state: success
real wall-clock time from trigger command to detected success (1s poll interval): 6.68 s200.8 ms against 6.68 s is a real, measured ~33x difference — and cityflow_hello is the smallest DAG in this course. validate_trips alone, inside cityflow_etl, is one task in a four-task chain that itself takes several real seconds; testing a single logic change to it by re-triggering the whole DAG and waiting for the scheduler’s own polling loop to notice is real, wasted time, every single time. None of that overhead — CLI trigger, scheduler heartbeat, worker pickup, polling interval — is present when a function is called directly in a Python process. This is the entire argument for pulling logic out into a plain function: not “it’s cleaner,” but a measured 33x.
Key Takeaways
cityflow_validation_logic.py’svalidate_trips_dfhas no Airflow import at all — it’s a plain function that takes a DataFrame and returns a dict, callable from a unit test, a notebook, or (Lesson 5) a real DAG task, without a scheduler running anywhere.- Real
assert-based tests should cover more than “one good case, one bad case” — this lesson’s edge cases (the exact100.0/100.01boundary, and a row failing two checks at once) caught behavior a single happy-path test would never have exercised. - A regression test against a previously-published real number (Module 5 Lesson 3’s
81013/69290/0.8553) is a strong way to prove an extracted function is faithful to the logic it replaced, not just superficially similar. - The 33x gap measured here (200.8 ms vs. 6.68 s) is real and specific to this machine and this DAG — the exact multiple will vary, but the direction never will: a scheduler, a worker, and a polling loop are real overhead a plain function call doesn’t pay.
Exercises
Exercise 1 — Add an eighth test for a row with a negative trip_distance (rather than one over the 100 mile ceiling) and confirm it’s still caught by trip_distance_out_of_range — the mask checks both directions (< 0 and > 100), but this lesson’s tests only ever exercised the upper bound directly.
Exercise 2 — Write a test for a DataFrame with zero rows (df_of() with no arguments). Confirm validate_trips_df doesn’t raise a ZeroDivisionError and instead returns pass_rate: 0.0 — check the if len(trips) else 0.0 guard in the source above and explain why it’s there.
Exercise 3 — Time the regression check (Test 7) on its own, separate from the other six tests, using time.perf_counter() around just that block. How much of the real 200.8 ms is the parquet read itself versus the five vectorized checks over 81,013 rows? Which one dominates?
Next lesson
Lesson 2 - Reading Logs