Skip to content
الصفحة الرئيسية
python
testing

From Fetch to Persist: Writing End‑to‑End Pytest for a Todo Pipeline

Mausam

When the first run of my nightly Airflow DAG silently dropped 27 % of the Todo rows, I realized I’d been treating the pipeline like a black box. Have you ever pushed a change, only to discover downstream reports are off by a few percent? In this post I walk you through the exact testing scaffolding that turned my “it works on my machine” pipeline into a safety‑net you can run on every commit. You’ll see how to unit‑test pure functions, mock the external API with requests‑mock, spin up an integration test that writes to a temporary filesystem, and add a lightweight benchmark so you know when a refactor slows you down.

Key Takeaways

  • Mocking the JSONPlaceholder endpoint isolates network failures and lets you verify request handling without flakiness.
  • Using tmp_path and caplog in pytest gives you a fast, filesystem‑level integration test that still exercises CSV persistence.
  • Parametrizing user‑ID scenarios uncovers edge cases like empty result sets and mixed‑completion states.
  • Embedding pytest‑benchmark in the suite surfaces performance regressions before they hit production.
  • Enforcing pytest‑cov thresholds in CI makes test coverage a gate‑keeper, not an after‑thought.

The Problem

My Airflow DAG pulls Todo items from JSONPlaceholder, filters out incomplete tasks, aggregates per‑user completion rates, and writes a CSV to an S3 bucket. The code works, but the test suite only covered the happy path of a single static response. Any change—new fields from the API, a regression in the aggregation logic, or a slowdown in the fetch step—could silently corrupt the nightly CSV, and I had no automated guard.

Data and Sources

Data accessed on 2026-10-31.

Step 1 — Design the Pipeline Skeleton

The pipeline consists of three pure functions (clean_todos, aggregate_stats, write_csv) and one orchestrator (run_pipeline). Keeping side‑effects isolated makes unit testing trivial.

def clean_todos(raw: list[dict]) -> list[dict]:
    """Keep only fields we care about and drop malformed rows."""
    cleaned = []
    for item in raw:
        if all(k in item for k in ("userId", "id", "title", "completed")):
            cleaned.append({
                "userId": int(item["userId"]),
                "id": int(item["id"]),
                "title": str(item["title"]).strip(),
                "completed": bool(item["completed"]),
            })
    return cleaned

Step 2 — Create Unit Tests for Pure Functions

Pure functions have no external dependencies, so a single pytest file suffices. Notice the use of pytest.mark.parametrize to cover multiple edge cases.

import pytest
from pipeline import clean_todos, aggregate_stats

@pytest.mark.parametrize("raw,expected_len", [
    ([], 0),
    ([{"userId":1,"id":1,"title":"a","completed":False}], 1),
    ([{"userId":1,"id":1}], 0),  # malformed row
])
def test_clean_todos(raw, expected_len):
    assert len(clean_todos(raw)) == expected_len

def test_aggregate_stats():
    sample = [
        {"userId": 1, "completed": True},
        {"userId": 1, "completed": False},
        {"userId": 2, "completed": True},
    ]
    stats = aggregate_stats(sample)
    assert stats[1]["completion_rate"] == 0.5
    assert stats[2]["completion_rate"] == 1.0

Step 3 — Mock the External API with requests‑mock

Instead of hitting the live endpoint, we register a fake response. This isolates network latency and lets us assert that the fetch logic raises on non‑200 status codes.

import requests
import requests_mock
from pipeline import fetch_todos

def test_fetch_success():
    with requests_mock.Mocker() as m:
        m.get("https://jsonplaceholder.typicode.com/todos", json=[{"userId":1,"id":1,"title":"x","completed":False}])
        data = fetch_todos()
        assert isinstance(data, list)
        assert data[0]["title"] == "x"

def test_fetch_failure():
    with requests_mock.Mocker() as m:
        m.get("https://jsonplaceholder.typicode.com/todos", status_code=500)
        with pytest.raises(RuntimeError):
            fetch_todos()

Step 4 — Use tmp_path and caplog for Integration Tests

An integration test spins up the whole pipeline, writes the CSV to a temporary directory, and inspects both the file contents and logged messages.

import csv
import logging

def test_end_to_end(tmp_path, caplog):
    # Mock the API
    with requests_mock.Mocker() as m:
        m.get("https://jsonplaceholder.typicode.com/todos", json=[
            {"userId":1,"id":1,"title":"Task A","completed":True},
            {"userId":1,"id":2,"title":"Task B","completed":False},
        ])
        output_file = tmp_path / "todos.csv"
        run_pipeline(output_path=str(output_file))
    
    # Verify CSV
    with open(output_file, newline="") as f:
        rows = list(csv.DictReader(f))
        assert rows[0]["userId"] == "1"
        assert rows[0]["completion_rate"] == "100.0"
    
    # Verify log
    assert any("Pipeline completed" in rec.message for rec in caplog.records)

Step 5 — Parameterize Tests for Different User Scenarios

Real‑world data contains users with zero todos, users with only completed tasks, and users with mixed states. Parametrization guarantees we hit each branch.

@pytest.mark.parametrize("api_response,expected_rate", [
    ([], None),  # no data → no CSV
    ([{"userId":3,"id":10,"title":"x","completed":True}], 100.0),
    ([{"userId":4,"id":11,"title":"y","completed":False}], 0.0),
])
def test_user_rate_variations(tmp_path, api_response, expected_rate):
    with requests_mock.Mocker() as m:
        m.get("https://jsonplaceholder.typicode.com/todos", json=api_response)
        out = tmp_path / "out.csv"
        run_pipeline(str(out))
        if expected_rate is None:
            assert not out.exists()
        else:
            with open(out) as f:
                rate = float(next(csv.DictReader(f))["completion_rate"])
                assert rate == expected_rate

Step 6 — Benchmark Performance with pytest‑benchmark

The fetch‑clean‑aggregate loop should stay under 200 ms on a typical CI runner. The benchmark fixture records the runtime and fails if it regresses beyond a threshold.

def test_performance(benchmark):
    result = benchmark(fetch_todos)
    # 0.2 seconds = 200 ms
    assert result < 0.2

Step 7 — Enforce Coverage and CI Integration

In pyproject.toml I added:

[tool.pytest.ini_options]
addopts = "--cov=pipeline --cov-fail-under=90 --benchmark-disable"

The --cov-fail-under flag aborts CI if coverage drops below 90 %, and --benchmark-disable prevents flaky CI runners from failing on the benchmark test.

The Problem

Summarizing, the core issue was a lack of confidence that the nightly DAG would keep producing correct CSVs as the code evolved. The solution is a layered pytest suite that validates each function in isolation, the end‑to‑end flow, and performance.

Loading the Data

The production fetch function handles HTTP errors and logs useful context.

import json
import logging
import requests

logger = logging.getLogger(__name__)

def fetch_todos(url: str = "https://jsonplaceholder.typicode.com/todos") -> list[dict]:
    try:
        resp = requests.get(url, timeout=5)
        resp.raise_for_status()
        return resp.json()
    except requests.RequestException as exc:
        logger.error("Failed to fetch todos: %s", exc)
        raise RuntimeError("Todo fetch failed") from exc

The Core Logic

After cleaning, we compute per‑user completion rates and write a CSV.

from collections import defaultdict
import csv
from pathlib import Path
from typing import List, Dict

def aggregate_stats(todos: List[Dict]) -> Dict[int, Dict]:
    stats = defaultdict(lambda: {"total": 0, "completed": 0})
    for t in todos:
        uid = t["userId"]
        stats[uid]["total"] += 1
        if t["completed"]:
            stats[uid]["completed"] += 1
    # Convert to completion rate percentage
    result = {}
    for uid, data in stats.items():
        rate = (data["completed"] / data["total"]) * 100 if data["total"] else 0
        result[uid] = {"completion_rate": round(rate, 2)}
    return result

def write_csv(stats: Dict[int, Dict], output_path: str) -> None:
    try:
        with open(output_path, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["userId", "completion_rate"])
            for uid, data in sorted(stats.items()):
                writer.writerow([uid, data["completion_rate"]])
        logger.info("Pipeline completed, CSV written to %s", output_path)
    except OSError as exc:
        logger.error("Failed to write CSV: %s", exc)
        raise

Putting It Together

The orchestrator wires the steps, catches top‑level errors, and returns an exit code suitable for Airflow’s PythonOperator.

def run_pipeline(output_path: str = "todos.csv") -> None:
    raw = fetch_todos()
    cleaned = clean_todos(raw)
    if not cleaned:
        logger.warning("No valid todos fetched; skipping CSV generation")
        return
    stats = aggregate_stats(cleaned)
    write_csv(stats, output_path)

if __name__ == "__main__":
    import sys
    try:
        run_pipeline()
    except Exception as e:
        logger.exception("Pipeline failed")
        sys.exit(1)

Complete Script

The full runnable script combining all steps:

#!/usr/bin/env python3
"""
End‑to‑end Todo pipeline with robust error handling.
Run: python pipeline.py
"""

import csv
import json
import logging
import sys
from collections import defaultdict
from pathlib import Path
from typing import List, Dict

import requests

# ----------------------------------------------------------------------
# Logging configuration
# ----------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)

# ----------------------------------------------------------------------
# Step 1: Fetch
# ----------------------------------------------------------------------
def fetch_todos(url: str = "https://jsonplaceholder.typicode.com/todos") -> List[Dict]:
    """Download the raw todo list, raising on network problems."""
    try:
        resp = requests.get(url, timeout=5)
        resp.raise_for_status()
        return resp.json()
    except requests.RequestException as exc:
        logger.error("Failed to fetch todos: %s", exc)
        raise RuntimeError("Todo fetch failed") from exc

# ----------------------------------------------------------------------
# Step 2: Clean
# ----------------------------------------------------------------------
def clean_todos(raw: List[Dict]) -> List[Dict]:
    """Validate required fields and coerce types."""
    cleaned = []
    for item in raw:
        if all(k in item for k in ("userId", "id", "title", "completed")):
            cleaned.append(
                {
                    "userId": int(item["userId"]),
                    "id": int(item["id"]),
                    "title": str(item["title"]).strip(),
                    "completed": bool(item["completed"]),
                }
            )
    return cleaned

# ----------------------------------------------------------------------
# Step 3: Aggregate
# ----------------------------------------------------------------------
def aggregate_stats(todos: List[Dict]) -> Dict[int, Dict]:
    """Return per‑user completion rate as a percentage."""
    stats = defaultdict(lambda: {"total": 0, "completed": 0})
    for t in todos:
        uid = t["userId"]
        stats[uid]["total"] += 1
        if t["completed"]:
            stats[uid]["completed"] += 1

    result = {}
    for uid, data in stats.items():
        rate = (data["completed"] / data["total"]) * 100 if data["total"] else 0
        result[uid] = {"completion_rate": round(rate, 2)}
    return result

# ----------------------------------------------------------------------
# Step 4: Persist
# ----------------------------------------------------------------------
def write_csv(stats: Dict[int, Dict], output_path: str) -> None:
    """Write the aggregation to a CSV file."""
    try:
        with open(output_path, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["userId", "completion_rate"])
            for uid, data in sorted(stats.items()):
                writer.writerow([uid, data["completion_rate"]])
        logger.info("Pipeline completed, CSV written to %s", output_path)
    except OSError as exc:
        logger.error("Failed to write CSV: %s", exc)
        raise

# ----------------------------------------------------------------------
# Orchestrator
# ----------------------------------------------------------------------
def run_pipeline(output_path: str = "todos.csv") -> None:
    raw = fetch_todos()
    cleaned = clean_todos(raw)
    if not cleaned:
        logger.warning("No valid todos fetched; skipping CSV generation")
        return
    stats = aggregate_stats(cleaned)
    write_csv(stats, output_path)

if __name__ == "__main__":
    try:
        run_pipeline()
    except Exception:
        logger.exception("Pipeline failed")
        sys.exit(1)

Expected Output

Running python pipeline.py creates todos.csv in the current directory. The first few lines look like:

userId,completion_rate
1,55.0
2,50.0
3,60.0
...

The console also prints an INFO line similar to:

2026-10-31 12:04:12 INFO __main__ - Pipeline completed, CSV written to todos.csv

Limitations and Tradeoffs

  • Mocked latency: requests‑mock removes real network time, so you might miss time‑outs that appear in production. Pair the mock suite with a nightly smoke test against the real endpoint.
  • Local CSV only: The script writes to the local filesystem; permission errors on S3 will not be caught. Add a dedicated test that injects a failing boto3 client or use moto for S3 mocks.
  • Benchmark variability: pytest‑benchmark is sensitive to CI hardware. In shared runners, disable it with --benchmark-disable and enable only on a dedicated performance node.
  • Parameter explosion: Over‑parameterizing user scenarios can make the suite slow and brittle. Limit combinations to those that surface distinct logic branches.

Frequently Asked Questions

Can I run the same test suite against a real S3 bucket?

Yes. Replace write_csv with a boto3 put_object call. In tests, use the moto library to mock S3 so you still get fast, deterministic runs.

What happens if the Todo API adds a new field tomorrow?

Because clean_todos only validates the four fields we need, extra keys are ignored automatically. If you later need the new field, extend the validation schema and add a unit test for the new path.

How do I make the Airflow DAG fail when the test suite breaks?

Configure your CI pipeline to publish a “latest‑successful‑run” tag to S3. The DAG’s first PythonOperator can read that tag; if the CI job failed, the tag won’t be updated and the operator raises AirflowException, halting the run.

Is there a way to generate documentation for the CSV schema automatically?

Annotate the output dict with a TypedDict and run pydoc-markdown or sphinx-autodoc. The generated markdown can be included in your project README.

Do I need to install requests‑mock in production?

No. It’s a test‑only dependency, listed under dev‑extras in pyproject.toml. Production environments only require requests, csv, and standard library modules.

What I'd Change

If I were to redesign this pipeline for a high‑throughput data lake, I’d swap the in‑memory aggregation for a streaming approach using polars or pandas with chunked reads, and I’d store the result as a Parquet file directly to S3. That would give me columnar compression, schema evolution, and native compatibility with downstream analytics tools. The pytest pattern stays the same—just replace the CSV writer with a Parquet writer and mock the S3 client—but the performance gains in production would be worth the extra complexity.

إرسال تعليق

Hi! How can we help you? Send us a message and we'll get back to you.