Testing DAGs and Wiring Up CI

Most broken DAGs never make it past a five-line import test — here's how to catch them before merge instead of at 3am.

The worst way to find out a DAG is broken is the scheduler telling you. A typo in an import, a task id that doesn’t match a downstream dependency, an accidental cycle — none of these need production to surface. They need a test that takes half a second, and a CI pipeline that runs it before anyone can merge. The good news is that Airflow DAGs are Python, and Python has a mature testing story. The trick is knowing which parts to test cheaply and which parts to test for real — and in Airflow 3.x, the tools for the “for real” half are better than the version most tutorials were written against.

One DAG to test against

Everything in this chapter tests the same small pipeline, so it’s worth putting on the table up front. It’s the bookshop’s daily load, boiled down to three TaskFlow tasks: pull orders from the API, normalize them, write them to the warehouse. The real logic lives in a plain function, normalize_orders, and the tasks are thin wrappers — which, as you’ll see, is the single decision that makes the whole thing testable.

# dags/load_bookshop.py
import pendulum
from airflow.sdk import dag, task
from include.bookshop import BookshopApiHook, BookshopWarehouseHook


def normalize_orders(rows: list[dict]) -> list[dict]:
    return [
        {
            "order_id": r["id"],
            "total_cents": int(round(r["total"] * 100)),
            "title": r["book"].strip(),
        }
        for r in rows
        if r["total"] > 0
    ]


@dag(
    dag_id="load_bookshop",
    schedule="@daily",
    start_date=pendulum.datetime(2026, 1, 1, tz="UTC"),
    catchup=False,
    params={"endpoint": "orders"},
)
def load_bookshop():
    @task
    def extract_orders(params: dict) -> list[dict]:
        return BookshopApiHook().fetch(params["endpoint"], since="{{ ds }}")

    @task
    def transform(rows: list[dict]) -> list[dict]:
        return normalize_orders(rows)

    @task
    def load_orders(rows: list[dict]) -> int:
        return BookshopWarehouseHook().load("stg_orders", rows)

    load_orders(transform(extract_orders()))


load_bookshop()

Three task ids — extract_orders, transform, load_orders — one param, and two hooks (the BookshopApiHook and BookshopWarehouseHook from the custom-operators chapter). That’s the whole system under test. Keep it in mind; every test below references this exact DAG.

Validation tests: the cheap net

Before you test what a DAG does, test that it’s a valid DAG at all. The single highest-value test in any Airflow repo just imports every DAG file and asserts nothing exploded. It catches the import typos, the missing dependencies, the malformed decorators — the boring breakage that’s most of what actually goes wrong.

import pytest
from airflow.models import DagBag


@pytest.fixture(scope="session")
def dagbag():
    return DagBag()


def test_no_import_errors(dagbag):
    assert not dagbag.import_errors, dagbag.import_errors


def test_expected_dags_present(dagbag):
    assert "load_bookshop" in dagbag.dags


def test_load_bookshop_structure(dagbag):
    dag = dagbag.dags["load_bookshop"]
    assert {t.task_id for t in dag.tasks} == {"extract_orders", "transform", "load_orders"}
    assert dag.get_task("transform").downstream_task_ids == {"load_orders"}


def test_bag_is_clean(dagbag):
    # No file blew up, and no DAG is a lone orphan with zero tasks.
    assert not dagbag.import_errors, dagbag.import_errors
    for dag_id, dag in dagbag.dags.items():
        assert dag.tasks, f"{dag_id} parsed with no tasks"
        assert dag.tags, f"{dag_id} has no tags — untag-able DAGs rot in the UI"

from airflow.models import DagBag is the one airflow.models import you should still reach for in 3.x — the rest of the modeling surface moved behind the Task SDK, but DagBag is the sanctioned way to parse a directory the way the scheduler would. Its import_errors dict is the same problem list you’d otherwise hit in production, surfaced in a unit test instead. Airflow builds the DAG graph at parse time, so a cycle won’t even load — test_no_import_errors catches it for free, no separate cycle assertion needed. Assert the task ids and the wiring, and you’ll notice the day someone renames transform and forgets to repoint load_orders at it. test_bag_is_clean widens the same idea to a whole-bag sweep: it walks every parsed DAG and asserts each one has tasks and carries at least one tag. The tag check is a small governance win — tags are how DAGs get filtered and owned in the 3.x UI, and an untagged DAG is one nobody finds until it fails. Because the bag is parsed once and shared, adding these blanket invariants costs nothing per DAG and scales to a repo of hundreds without a per-DAG test each.

DagBag has a couple of sharp edges worth knowing before you lean on it:

  • Forget include_examples=False. You’ll see it in every 2.x example, and in Airflow 3 it is a TypeError — the kwarg was removed. You also don’t need it: a bare DagBag() reads your DAGs folder and nothing else. If your project keeps DAGs somewhere non-standard, pass the path: DagBag(dag_folder="dags/").
  • The whole-bag parse can time out. DagBag enforces core.dagbag_import_timeout per file — a DAG that does real work at import time (hits a network, reads a big file) will blow the timeout in CI even though it “works” locally. That’s a signal, not a flake: move that work inside a task. If you only care about one file, parse one file — DagBag(dag_folder="dags/load_bookshop.py") is faster and gives a tighter failure than parsing the whole directory.
  • Session scope pays off. Parsing is the expensive part, so scope the fixture to the session and let every test share one parsed bag.

Unit tests: pull the logic out of Airflow

Here’s the mistake that makes DAGs feel untestable: burying real logic inside a task. The moment your transformation lives in the body of an @task function, testing it means standing up a runtime. So don’t do that — which is exactly why transform in the canonical DAG is a one-liner that calls normalize_orders. The task is the wrapper; the function is the thing you test.

from dags.load_bookshop import normalize_orders


def test_normalize_drops_zero_total_orders():
    rows = [{"id": 1, "total": 12.5, "book": " Dune "},
            {"id": 2, "total": 0, "book": "Free Sample"}]
    out = normalize_orders(rows)
    assert len(out) == 1
    assert out[0] == {"order_id": 1, "total_cents": 1250, "title": "Dune"}

That’s the test that earns its keep. It pins the actual business rule — drop zero-total orders, round to cents, trim titles — and runs in microseconds, with no Airflow in sight. You can throw fuzzy inputs at it, table-drive it, cover the rounding edge cases, all without a scheduler. Every rule you can express as a plain function is a rule you can test this way, and most of them can be.

dag.test(): run the whole DAG in-process

Unit tests prove the logic. They don’t prove the wiring — that extract_orders hands the right shape to transform, that transform feeds load_orders, that the XCom plumbing between TaskFlow tasks actually connects. For that you want to run the DAG, and in Airflow 3.x you run it with dag.test().

def test_load_bookshop_runs(dagbag):
    dag = dagbag.dags["load_bookshop"]
    dr = dag.test()
    assert dr.state == "success"

dag.test() executes the entire DAG in your current process — no scheduler, no worker, no Celery, no separate triggerer. It runs the tasks in dependency order, moves real XCom between them, and returns the completed DagRun so you can assert on its state. A task that raises surfaces the traceback right there in your test output, which is exactly what you want from a test.

It is the programmatic sibling of the airflow dags test <dag_id> CLI command — which, contrary to a lot of migration chatter, is not gone in 3.x; it’s alive and has even gained flags (--bundle-name, --dagfile-path, --mark-success-pattern). Use the CLI from a shell, dag.test() from a test or a debugger. If you have a shell script or a runbook that shells out to airflow dags test, it no longer works — the in-process runner is now a method on the DAG object, called from Python. That’s a strict upgrade for testing: dag.test() lives in your test file, takes fixtures, and composes with monkeypatch and pytest the way the CLI never could.

It runs real task code, so give it real dependencies. This is the part people trip on. dag.test() doesn’t mock anything — when extract_orders calls BookshopApiHook().fetch(...), that call really happens. Run the canonical DAG untouched and it’ll try to reach the bookshop API. So the pattern is: patch the boundaries, run the middle for real.

from include import bookshop


def test_load_bookshop_end_to_end(dagbag, monkeypatch):
    monkeypatch.setattr(
        bookshop.BookshopApiHook, "fetch",
        lambda self, endpoint, since: [{"id": 1, "total": 12.5, "book": " Dune "}],
    )
    captured = {}
    monkeypatch.setattr(
        bookshop.BookshopWarehouseHook, "load",
        lambda self, table, rows: captured.setdefault(table, rows) or len(rows),
    )

    dr = dagbag.dags["load_bookshop"].test()

    assert dr.state == "success"
    assert captured["stg_orders"][0]["title"] == "Dune"

Now the run is hermetic: the API hook returns a fixture, the warehouse hook captures what it was asked to write, and everything between — the extract-transform-load wiring, the normalize_orders call, the XCom hand-offs — runs for real. The assertion proves the pipeline end to end: a raw order went in one side and a trimmed, normalized row came out the other.

dag.test() takes a few arguments that matter for realistic tests (exact keyword names are docs-verified against 3.3.0, not run here):

  • run_conf= feeds the run’s conf, which is how params get overridden. The canonical DAG defaults endpoint to "orders"; a test can drive the returns path with dag.test(run_conf={"endpoint": "orders/returns"}) and assert the hook was called with "orders/returns".
  • conn_file_path= and variable_file_path= point at a YAML/JSON file of connections and variables to load for the run, so you can supply bookshop_api credentials or a warehouse_schema variable without touching your real metadata DB. Handy when a task genuinely needs a connection to exist.
  • logical_date= sets the run’s date, which is what {{ ds }} in extract_orders renders from — pass it when a test depends on a specific partition date.
  • use_executor=False is the default and the right one for tests: it runs tasks inline. Flip it on only if you specifically want to exercise a real executor path.

Know its limits, too. dag.test() runs tasks sequentially in one process, so it won’t catch a race between two parallel tasks, it doesn’t exercise pools or concurrency limits, and it won’t reproduce a scheduler-timing bug. Deferrable operators don’t actually defer to a triggerer under it — they run through inline. It’s a faithful test of logic and wiring, not of scheduling behavior. For the scheduling half, you’re back to a real deployment; for everything else, dag.test() is the fastest honest test you have.

Testing custom operators and hooks

The custom BookshopApiToWarehouseOperator from earlier in this series is a class, and classes test even more directly than DAGs — you don’t need dag.test() at all. An operator’s contract is its execute method, so call it with a context and assert on what it returns and what it did.

from include import bookshop
from include.bookshop import BookshopApiToWarehouseOperator


def test_operator_execute(monkeypatch):
    op = BookshopApiToWarehouseOperator(
        task_id="load", conn_id="bookshop_api",
        endpoint="orders", logical_date="2026-07-01", table="stg_orders",
    )
    monkeypatch.setattr(bookshop.BookshopApiHook, "fetch",
                        lambda self, endpoint, since: [{"id": 1}, {"id": 2}])
    monkeypatch.setattr(bookshop.BookshopApiHook, "load_to_warehouse",
                        lambda self, rows, table: None)

    assert op.execute(context={}) == 2

Instantiate the operator, mock the hook methods it calls, hand execute a plain dict for the context, and assert. No DAG, no runtime. The context can be {} when the operator doesn’t read from it, or a fake with just the keys it does read — more on that below.

The other thing worth testing on an operator is templating, because template_fields is easy to get wrong and silent when you do. render_template_fields runs the same Jinja pass the scheduler would, against a context you supply:

def test_operator_renders_templates():
    op = BookshopApiToWarehouseOperator(
        task_id="load", conn_id="bookshop_api",
        endpoint="{{ params.endpoint }}", logical_date="{{ ds }}", table="stg_orders",
    )
    op.render_template_fields({"params": {"endpoint": "returns"}, "ds": "2026-07-01"})

    assert op.endpoint == "returns"
    assert op.logical_date == "2026-07-01"

If someone forgets to list endpoint in template_fields, this test fails with op.endpoint == "{{ params.endpoint }}" — the literal string, unrendered — which is precisely the bug that ships to production and surprises you a week later.

When you genuinely need a bound TaskInstance — to test something that reads ti.try_number or the map index — Airflow ships pytest fixtures (dag_maker, create_task_instance) for constructing one. In 3.x, though, a live TaskInstance runs through the Task Execution API, so building one standalone is more ceremony than it used to be; reach for those fixtures only when .execute(context) genuinely can’t express the test, and prefer dag.test() for anything that needs a real run.

Hooks are the easiest of all, because a hook is a plain object whose only Airflow dependency is get_connection. Mock that and the whole thing is unit-testable:

from types import SimpleNamespace
from include import bookshop


def test_hook_builds_authed_session(monkeypatch):
    fake_conn = SimpleNamespace(host="https://api.bookshop.test", password="tok_123")
    monkeypatch.setattr(bookshop.BookshopApiHook, "get_connection",
                        lambda self, conn_id: fake_conn)

    session = bookshop.BookshopApiHook().get_conn()

    assert session.headers["Authorization"] == "Bearer tok_123"
    assert session.base_url == "https://api.bookshop.test"

No connection in a metadata DB, no environment variable — just a fake connection object handed back from get_connection, and an assertion that the hook wired the auth header and base URL correctly. This is the payoff of keeping the integration in a hook: the trickiest code in your project (auth, pagination, retries) becomes ordinary Python you can cover without Airflow running at all.

Faking the context, mocking XCom

Some operators read from the context and push or pull XCom directly. You don’t need a real one. The context is a dict-like object, so a fake with the handful of keys the code touches is enough, and a MagicMock task-instance handles XCom:

from unittest.mock import MagicMock


def test_operator_fetches_and_loads(monkeypatch):
    fetched = [{"id": 1, "total": 12.5, "book": " Dune "}]
    loaded = {}

    monkeypatch.setattr(BookshopApiHook, "fetch",
                        lambda self, endpoint, since: fetched)
    monkeypatch.setattr(BookshopApiHook, "load_to_warehouse",
                        lambda self, rows, table: loaded.update(rows=rows, table=table))

    op = BookshopApiToWarehouseOperator(
        task_id="load", conn_id="bookshop_api_default",
        endpoint="orders", since="2026-07-01", table="stg_orders",
    )
    count = op.execute(context={})

    assert count == 1                       # execute's return value becomes the XCom
    assert loaded["table"] == "stg_orders"

That’s BookshopApiToWarehouseOperator from chapter 01 — the operator you built — with its hook stubbed out. You assert on what it asked the hook for and what it handed back. Note there is no xcom_push to assert on: execute simply returns the count, and Airflow turns the return value into the XCom for you. That is the TaskFlow contract, and it means the thing to test is the return value, not a mock’s call log. For connections and variables specifically, there’s an even lighter trick than mocking: environment overrides. Airflow reads a connection named bookshop_api from AIRFLOW_CONN_BOOKSHOP_API and a variable warehouse_schema from AIRFLOW_VAR_WAREHOUSE_SCHEMA when they’re set, so monkeypatch.setenv is often the cleanest way to satisfy code that calls BaseHook.get_connection or Variable.get:

def test_reads_connection_from_env(monkeypatch):
    monkeypatch.setenv("AIRFLOW_CONN_BOOKSHOP_API",
                       "https://[email protected]")
    monkeypatch.setenv("AIRFLOW_VAR_WAREHOUSE_SCHEMA", "test_raw")
    # code under test can now call get_connection("bookshop_api")
    # and Variable.get("warehouse_schema") with no metadata DB

Data quality vs. pipeline testing — and the tricky constructs

There are two different things people mean by “testing a pipeline,” and it’s worth being precise. Pipeline testing is everything above: does the code do what I think, in CI, before merge. Data-quality testing is a runtime concern — does today’s data satisfy its contract — and it belongs inside the DAG, as tasks or expectations that run against the real load, not in your pytest suite. A CI test asserting “row counts are positive” tests nothing, because there’s no data in CI. Keep the two separate: pytest proves the pipeline; a data-quality task (dbt tests, Great Expectations, a SQLColumnCheckOperator) proves the run.

The constructs from earlier chapters — mapping, branching, trigger rules — each have a natural pipeline test, and they’re the ones most projects skip:

  • Mapping. Test the fan-out width by driving the mapped task through dag.test() with a fixture list and asserting how many instances ran, or — cheaper — test the function that produces the list. The bug you’re guarding against is “the upstream returned something that isn’t a list,” which a plain assert isinstance(list_region_files(...), list) catches before mapping ever explodes at runtime. When the fan-out width itself is the contract — “one task per region file, and no more” — assert it against the real run. dag.test() returns the DagRun, and its task instances carry a map_index, so you can count the expansion and even check per-index output:

    from include import bookshop

def test_per_region_load_fans_out(dagbag, monkeypatch):
    # Upstream lists three region files; the mapped loader runs once per file.
    monkeypatch.setattr(
        bookshop.BookshopStorageHook, "list_keys",
        lambda self, prefix: ["us.json", "eu.json", "apac.json"],
    )
    monkeypatch.setattr(
        bookshop.BookshopWarehouseHook, "copy_into",
        lambda self, key, table, batch_rows: {
            "us.json": 12, "eu.json": 7, "apac.json": 3
        }[key],
    )

    dr = dagbag.dags["bookshop_regional_load"].test()

    mapped = [ti for ti in dr.get_task_instances()
              if ti.task_id == "load_region" and ti.map_index >= 0]
    assert len(mapped) == 3                       # fan-out width
    assert all(ti.state == "success" for ti in mapped)
    assert sorted(ti.map_index for ti in mapped) == [0, 1, 2]

    That’s bookshop_regional_load from chapter 03, with its two hooks stubbed. The ti.map_index >= 0 filter matters: an unexpanded task carries map_index == -1, and the reduce/summary side of a mapped group can show up separately, so you count only the real expanded instances. This is the test that catches the day someone changes list_keys to return a dict or a generator — the width silently drops to one, and only this assertion notices.

  • Branching. A branch is a function that returns a task id (or ids). Test it as a function — no DAG required:

    from dags.bookshop_conditional import choose_path

def test_branch_picks_the_load_path():
    assert choose_path.function(128) == "load_orders"
    assert choose_path.function(0) == "skip_load"

    Note .function — a @task.branch-decorated object is a task, not a plain callable, so you reach through to the undecorated function to call it directly. That is the whole trick to testing branches: the decision is ordinary Python, and ordinary Python is cheap to test.

  • Trigger rules. The classic branching bug is a join task that never runs because upstream branches got skipped rather than succeeded. Test it structurally — assert the join’s trigger rule is the tolerant one — and then, if it’s load-bearing, exercise the skip path through dag.test() and assert the join still succeeded:

    def test_join_tolerates_skipped_branches(dagbag):
    join = dagbag.dags["bookshop_conditional"].get_task("report")
    assert join.trigger_rule == "none_failed_min_one_success"

That last assertion catches the single most common conditional-DAG mistake — a join left on the default all_success rule, which quietly skips itself the moment one branch is skipped — in a test that runs in a millisecond.

CI: run it all on every PR

Local tests you have to remember to run aren’t a safety net. Wire them into CI so a broken DAG can’t merge. Here’s a GitHub Actions workflow that does the real work: it runs the suite against a matrix of Airflow versions, caches the environment so it’s fast, stands up a real Postgres for the metadata DB, and lints for the 2.x-isms that 3.x deleted.

name: ci
on: pull_request

jobs:
  test:
    runs-on: ubuntu-latest
    strategy:
      fail-fast: false
      matrix:
        airflow: ["3.2.0", "3.3.0"]
    services:
      postgres:
        image: postgres:16
        env:
          POSTGRES_USER: airflow
          POSTGRES_PASSWORD: airflow
          POSTGRES_DB: airflow
        ports: ["5432:5432"]
        options: >-
          --health-cmd pg_isready --health-interval 5s
          --health-timeout 5s --health-retries 5
    env:
      AIRFLOW__DATABASE__SQL_ALCHEMY_CONN: postgresql+psycopg2://airflow:airflow@localhost:5432/airflow
      AIRFLOW__CORE__LOAD_EXAMPLES: "False"
    steps:
      - uses: actions/checkout@v4
      - uses: astral-sh/setup-uv@v5
        with:
          enable-cache: true
      - run: uv pip install --system "apache-airflow==${{ matrix.airflow }}" -r requirements.txt pytest ruff
      - run: airflow db migrate
      - run: ruff check --select AIR30 dags/ include/
      - run: pytest tests/ -q

Each part pulls its weight. The matrix runs the suite against the Airflow you’re on today and the one you’ll upgrade to next, so an API that 3.3 deprecated shows up as a failing cell before the upgrade, not during it. setup-uv with enable-cache: true caches the resolved environment keyed on your lockfile, so the second run installs in seconds — the same caching you’d get from actions/cache around a pip download dir, with less wiring. The Postgres service container gives the metadata DB a real home: airflow db migrate creates the schema against it, so dag.test() runs and any test that touches the DB behaves like production instead of against an ephemeral SQLite that hides locking bugs.

ruff check --select AIR30 is the line that earns its place in an Airflow 3.x repo specifically. The AIR30x rule families flag exactly the removed-and-renamed surface this whole series keeps warning about — schedule_interval=, timetable=, sla=, SubDagOperator, the old airflow.operators.* and airflow.models import paths — so the stale 2.x snippet someone copied off a three-year-old blog post fails the build with a clear message instead of a mysterious parse error six weeks later. Add it to pre-commit and it runs before the commit even lands:

# .pre-commit-config.yaml
repos:
  - repo: https://github.com/astral-sh/ruff-pre-commit
    rev: v0.9.0
    hooks:
      - id: ruff
        args: ["--select", "AIR30", "--fix"]
  - repo: https://github.com/gitleaks/gitleaks
    rev: v8.24.0
    hooks:
      - id: gitleaks

The second hook matters more than it looks. DAG files are where credentials go to die — someone hardcodes a warehouse password “just to test,” it lands in git history, and now it’s a rotation. A secret-scanner in pre-commit (and again in CI) refuses the commit that carries a bearer token or a connection string, which is the one CI failure your security team will thank you for.

If your project is an Astro (Astronomer) one, there’s a nicer front door than raw pytest: astro dev pytest runs your test suite inside the project’s own Docker image — the same image astro dev start runs — so the tests execute against the exact Airflow version, providers, and system libraries your deployment uses. Wiring astro dev pytest into CI (instead of, or alongside, a bare pytest) closes the “works on my Airflow” gap entirely, because there’s only one Airflow: the one in the image. On merge to main, a second workflow deploys — which is where the next post picks up.

Final thoughts

The reason Airflow testing has a bad reputation is that people try to test the framework instead of their code. Airflow’s scheduler works; you don’t need to prove it. What you need to prove is that your transform does the right thing, your operator renders its templates, and your DAG is wired the way you think — and in 3.x every one of those is an ordinary Python test: a plain function for the logic, .execute(context) for the operator, dag.test() for the wiring, a five-line DagBag import for the cheap net. Write that import test today, before anything else. It’s the highest ratio of bugs-caught to effort you’ll find in this whole series.

Next: Shipping It: Astronomer, MWAA, and Cloud Composer

Comments