Skip to content
Codeloom
Airflow

Testing Airflow DAGs: A Comprehensive Guide

Unit test tasks, validate DAG structure, mock external services, and build CI/CD pipelines for Airflow. Prevent silent data loss with systematic testing.

·13 min read · By Codeloom
Intermediate 18 min read

What you'll learn

  • Why untested DAGs are a ticking time bomb for data quality
  • DAG validation tests that catch import errors and structural problems
  • Unit testing individual tasks with pytest
  • Mocking external services like S3, BigQuery, and APIs
  • Building a CI/CD pipeline that prevents broken DAGs from reaching production

Prerequisites

  • Experience writing Airflow DAGs
  • Basic familiarity with pytest
  • Understanding of Python imports and modules

Why Testing Airflow DAGs Matters

Here is how most Airflow failures play out. A developer makes a small change to a DAG — updates a SQL query, renames a connection, adds a new task. They test it manually by triggering a run in their local Airflow instance. It works. They push it to production. Three days later, someone notices that Tuesday’s revenue numbers are wrong. The investigation reveals that the modified SQL query had a subtle bug that produced incorrect aggregations, but only for certain date ranges. The DAG ran successfully every night — no errors, no alerts, just silently wrong data feeding into dashboards and reports.

This is the fundamental danger of untested DAGs: broken DAGs do not always fail loudly. They fail silently, producing data that looks right but is not. By the time you discover the problem, the damage has compounded — downstream reports are wrong, decisions have been made on bad data, and you need to backfill and reconcile days or weeks of results.

Testing is the practice that prevents this. Not all testing — you do not need 100% code coverage on every DAG. But a small, targeted set of tests catches the vast majority of problems before they reach production.

Layer 1: DAG Validation Tests

The highest-value, lowest-effort tests you can write are DAG validation tests. These tests load your DAG files and check for problems that would prevent them from running at all. Think of them as a smoke test — they do not verify business logic, but they catch structural failures that would otherwise only surface when the scheduler tries to parse your files.

The Import Test

The single most valuable test in any Airflow project is the import test. It loads every DAG file through Airflow’s DagBag and checks that none of them produce import errors. This catches syntax errors, missing Python packages, broken imports, and any module-level code that throws exceptions.

# tests/test_dag_validation.py
import pytest
from airflow.models import DagBag

@pytest.fixture(scope="session")
def dag_bag():
    """Load all DAGs once for the test session."""
    return DagBag(dag_folder="dags/", include_examples=False)

def test_no_import_errors(dag_bag):
    """Every DAG file must import without errors."""
    assert len(dag_bag.import_errors) == 0, (
        f"DAG import errors found:\n"
        + "\n".join(
            f"  {path}: {err}"
            for path, err in dag_bag.import_errors.items()
        )
    )

This test takes seconds to run and prevents a surprising number of production incidents. Common catches include:

  • A developer removes a utility module that a DAG imports
  • A new dependency is added to a DAG but not to requirements.txt
  • A YAML config file that drives dynamic DAG generation has a syntax error
  • A Python 3.10+ syntax feature is used but production runs Python 3.9

The “No Cycles” Test

Airflow DAGs must be acyclic (the “A” in DAG). If task A depends on task B which depends on task A, the DAG is invalid and will not schedule. While Airflow catches this at parse time, an explicit test makes the error message clearer and catches it earlier in the development cycle.

def test_no_cycles(dag_bag):
    """No DAG should contain circular dependencies."""
    for dag_id, dag in dag_bag.dags.items():
        # Airflow's DagBag will not import DAGs with cycles,
        # but we verify no partial-load issues
        assert dag is not None, f"DAG {dag_id} failed to load"

Structure Tests for Critical DAGs

For your most important pipelines — the ones that feed executive dashboards, power customer-facing features, or handle financial data — add tests that verify the expected structure. These tests act as a contract: if someone accidentally removes a task or changes a dependency, the test fails before the change reaches production.

def test_revenue_pipeline_structure(dag_bag):
    """The revenue pipeline must have the expected tasks and dependencies."""
    dag = dag_bag.get_dag("daily_revenue")
    assert dag is not None, "daily_revenue DAG not found"

    task_ids = {t.task_id for t in dag.tasks}

    # Required tasks must exist
    required = {"extract_orders", "extract_refunds", "transform", "load", "notify"}
    missing = required - task_ids
    assert not missing, f"Missing tasks: {missing}"

    # Verify critical dependency: transform must run after both extracts
    transform = dag.get_task("transform")
    upstream_ids = {t.task_id for t in transform.upstream_list}
    assert "extract_orders" in upstream_ids
    assert "extract_refunds" in upstream_ids

def test_all_dags_have_owners(dag_bag):
    """Every DAG must have a non-default owner."""
    for dag_id, dag in dag_bag.dags.items():
        owner = dag.default_args.get("owner", "airflow")
        assert owner != "airflow", (
            f"{dag_id} uses default owner 'airflow' -- set a team owner"
        )

def test_all_dags_have_tags(dag_bag):
    """Every DAG must have at least one tag for filtering in the UI."""
    for dag_id, dag in dag_bag.dags.items():
        assert dag.tags, f"{dag_id} has no tags"

Configuration Tests

If your DAGs use configuration values (schedules, connection IDs, pool names), test that they are set to valid values:

VALID_SCHEDULES = {"@daily", "@hourly", "@weekly", None}

def test_schedules_are_valid(dag_bag):
    """Catch typos in schedule expressions."""
    for dag_id, dag in dag_bag.dags.items():
        schedule = dag.schedule_interval
        if isinstance(schedule, str) and not schedule.startswith("@"):
            # Cron expression -- basic format check
            parts = schedule.split()
            assert len(parts) == 5, (
                f"{dag_id} has invalid cron: '{schedule}'"
            )

Layer 2: Unit Testing Task Logic

DAG validation tests verify structure. Unit tests verify behavior — that your tasks produce the right output given specific input. The key principle here is separation of concerns: your most testable DAG code is code where the business logic lives in plain Python functions that have no Airflow dependencies.

Extract Business Logic into Testable Functions

Instead of writing all your logic inside @task-decorated functions, extract the core logic into plain Python functions in a utility module. These functions take inputs and return outputs — no Airflow context, no hooks, no operators. Test them with standard pytest.

# dags/utils/transforms.py
from typing import Optional
from datetime import datetime

def normalize_email(email: str) -> str:
    """Lowercase and strip whitespace from email."""
    return email.strip().lower()

def calculate_revenue(
    orders: list[dict],
    refunds: list[dict],
) -> float:
    """Net revenue = gross orders minus refunds."""
    gross = sum(o["amount"] for o in orders)
    refunded = sum(r["amount"] for r in refunds)
    return gross - refunded

def partition_key(ds: str) -> str:
    """Convert Airflow ds (YYYY-MM-DD) to partition format."""
    dt = datetime.strptime(ds, "%Y-%m-%d")
    return dt.strftime("%Y/%m/%d")

def validate_record(record: dict) -> Optional[dict]:
    """Return the record if valid, None if it should be filtered."""
    if not record.get("email"):
        return None
    if record.get("amount", 0) < 0:
        return None
    record["email"] = normalize_email(record["email"])
    return record
# tests/test_transforms.py
import pytest
from dags.utils.transforms import (
    normalize_email,
    calculate_revenue,
    partition_key,
    validate_record,
)

class TestNormalizeEmail:
    def test_strips_whitespace(self):
        assert normalize_email("  user@test.com  ") == "user@test.com"

    def test_lowercases(self):
        assert normalize_email("User@Example.COM") == "user@example.com"

class TestCalculateRevenue:
    def test_no_refunds(self):
        orders = [{"amount": 100}, {"amount": 200}]
        assert calculate_revenue(orders, []) == 300

    def test_with_refunds(self):
        orders = [{"amount": 100}, {"amount": 200}]
        refunds = [{"amount": 50}]
        assert calculate_revenue(orders, refunds) == 250

    def test_empty_orders(self):
        assert calculate_revenue([], []) == 0

class TestPartitionKey:
    def test_formats_correctly(self):
        assert partition_key("2026-07-12") == "2026/07/12"

class TestValidateRecord:
    def test_valid_record(self):
        r = {"email": "a@b.com", "amount": 10}
        assert validate_record(r) is not None

    def test_missing_email_returns_none(self):
        assert validate_record({"amount": 10}) is None

    def test_negative_amount_returns_none(self):
        assert validate_record({"email": "a@b.com", "amount": -5}) is None

These tests run fast (no Airflow, no database, no network calls), are easy to write, and catch the bugs that matter most — incorrect business logic.

Testing Tasks That Use Airflow Context

Some tasks need Airflow context (the execution date, the DAG run ID, configuration). You can test these by constructing a mock context dict:

# dags/tasks/extract.py
from airflow.decorators import task
from airflow.providers.postgres.hooks.postgres import PostgresHook

@task
def extract_daily_orders(**context):
    ds = context["ds"]
    hook = PostgresHook("orders_db")
    query = f"SELECT * FROM orders WHERE order_date = '{ds}'"
    return hook.get_records(query)
# tests/test_extract.py
from unittest.mock import patch, MagicMock

def test_extract_builds_correct_query():
    """Verify the query uses the execution date correctly."""
    mock_hook = MagicMock()
    mock_hook.get_records.return_value = [("order_1", 100)]

    with patch(
        "dags.tasks.extract.PostgresHook", return_value=mock_hook
    ):
        from dags.tasks.extract import extract_daily_orders

        context = {"ds": "2026-07-12"}
        # Call the underlying function (not the decorated version)
        result = extract_daily_orders.function(**context)

    mock_hook.get_records.assert_called_once_with(
        "SELECT * FROM orders WHERE order_date = '2026-07-12'"
    )
    assert result == [("order_1", 100)]

The .function attribute on a @task-decorated function gives you access to the underlying Python function, bypassing the Airflow task machinery. This lets you call it directly in tests.

Layer 3: Mocking External Services

Most Airflow tasks interact with external systems — databases, cloud storage, APIs. Testing these interactions without actually hitting the real services requires mocking.

Mocking S3 S3 with moto

The moto library provides mock implementations of AWS services. It is the standard approach for testing code that interacts with S3, SQS, DynamoDB, and other AWS services.

# tests/test_s3_upload.py
import boto3
import pytest
from moto import mock_aws

@mock_aws
def test_upload_to_s3():
    """Test that our upload function creates the expected S3 object."""
    # Create a mock S3 bucket
    s3 = boto3.client("s3", region_name="us-east-1")
    s3.create_bucket(Bucket="data-lake")

    # Call your function
    from dags.utils.s3_helpers import upload_csv
    upload_csv(
        bucket="data-lake",
        key="2026/07/12/orders.csv",
        data=[{"id": 1, "amount": 100}],
    )

    # Verify the object was created
    obj = s3.get_object(Bucket="data-lake", Key="2026/07/12/orders.csv")
    content = obj["Body"].read().decode()
    assert "id,amount" in content
    assert "1,100" in content

Mocking BigQuery BigQuery

For BigQuery, the approach depends on what you are testing. If you are testing query construction, mock the client. If you are testing query results, consider using a local database as a stand-in.

# tests/test_bigquery_load.py
from unittest.mock import patch, MagicMock

def test_bigquery_load_calls_correct_table():
    mock_hook = MagicMock()

    with patch(
        "dags.tasks.load.BigQueryHook", return_value=mock_hook
    ):
        from dags.tasks.load import load_to_bq

        load_to_bq.function(
            records=[{"user": "alice", "score": 95}],
            destination_table="analytics.user_scores",
            ds="2026-07-12",
        )

    # Verify the hook was called with expected params
    mock_hook.insert_rows.assert_called_once()
    call_args = mock_hook.insert_rows.call_args
    assert call_args[0][0] == "analytics.user_scores"

Mocking HTTP APIs

For tasks that call external APIs, use responses or requests-mock to intercept HTTP calls:

# tests/test_api_extract.py
import responses

@responses.activate
def test_extract_from_api():
    """Test API extraction with a mocked response."""
    responses.add(
        responses.GET,
        "https://api.vendor.com/v1/events",
        json={"events": [{"id": 1, "type": "click"}]},
        status=200,
    )

    from dags.utils.api_client import fetch_events
    result = fetch_events(date="2026-07-12")

    assert len(result) == 1
    assert result[0]["type"] == "click"

@responses.activate
def test_extract_handles_rate_limit():
    """Test that our client handles 429 responses gracefully."""
    responses.add(
        responses.GET,
        "https://api.vendor.com/v1/events",
        json={"error": "rate limited"},
        status=429,
    )

    from dags.utils.api_client import fetch_events
    with pytest.raises(RateLimitError):
        fetch_events(date="2026-07-12")

Layer 4: Integration Testing

Unit tests verify individual pieces. Integration tests verify that the pieces work together. For Airflow, this means running tasks against real (but disposable) databases and services.

Using a Test Database

Spin up a PostgreSQL container for your test suite using Docker, and run your extract and load logic against it:

# conftest.py
import pytest
import subprocess

@pytest.fixture(scope="session")
def test_db():
    """Start a PostgreSQL container for integration tests."""
    subprocess.run([
        "docker", "run", "-d",
        "--name", "airflow-test-db",
        "-e", "POSTGRES_PASSWORD=test",
        "-p", "5433:5432",
        "postgres:15",
    ], check=True)

    # Wait for the database to be ready
    import time
    time.sleep(3)

    yield {
        "host": "localhost",
        "port": 5433,
        "user": "postgres",
        "password": "test",
        "database": "postgres",
    }

    subprocess.run(
        ["docker", "rm", "-f", "airflow-test-db"], check=True
    )
# tests/test_integration.py
import psycopg2

def test_load_is_idempotent(test_db):
    """Running load twice produces the same result."""
    conn = psycopg2.connect(**test_db)
    cur = conn.cursor()

    cur.execute("""
        CREATE TABLE IF NOT EXISTS events (
            id INT, name TEXT, partition_date DATE
        )
    """)
    conn.commit()

    from dags.utils.loader import load_events
    records = [(1, "click", "2026-07-12"), (2, "view", "2026-07-12")]

    # Load twice
    load_events(conn, records, "2026-07-12")
    load_events(conn, records, "2026-07-12")

    cur.execute("SELECT COUNT(*) FROM events WHERE partition_date = '2026-07-12'")
    count = cur.fetchone()[0]
    assert count == 2, f"Expected 2 rows (idempotent), got {count}"

    conn.close()

Building a CI/CD Pipeline for Airflow DAGs

Tests are only valuable if they run automatically. A CI/CD pipeline ensures that every change to your DAGs goes through validation, unit tests, and optionally integration tests before reaching production.

GitHub Actions GitHub Actions Pipeline

# .github/workflows/airflow-ci.yml
name: Airflow DAG CI

on:
  pull_request:
    paths:
      - "dags/**"
      - "tests/**"
      - "requirements.txt"

jobs:
  validate-and-test:
    runs-on: ubuntu-latest

    services:
      postgres:
        image: postgres:15
        env:
          POSTGRES_PASSWORD: test
        ports:
          - 5433:5432
        options: >-
          --health-cmd pg_isready
          --health-interval 10s
          --health-timeout 5s
          --health-retries 5

    steps:
      - uses: actions/checkout@v4

      - uses: actions/setup-python@v5
        with:
          python-version: "3.11"

      - name: Install dependencies
        run: |
          pip install apache-airflow==2.9.0 pytest pytest-cov ruff
          pip install -r requirements.txt

      - name: Lint DAG files
        run: ruff check dags/

      - name: Run DAG validation tests
        run: pytest tests/test_dag_validation.py -v

      - name: Run unit tests
        run: pytest tests/ -v --ignore=tests/test_integration.py

      - name: Run integration tests
        env:
          TEST_DB_HOST: localhost
          TEST_DB_PORT: 5433
          TEST_DB_PASSWORD: test
        run: pytest tests/test_integration.py -v

What the Pipeline Catches

This pipeline prevents several categories of failures:

CheckWhat it catches
ruff checkStyle violations, unused imports, obvious bugs
Import testBroken imports, missing packages, syntax errors
Structure testsAccidentally removed tasks, broken dependencies
Unit testsIncorrect business logic, wrong transformations
Integration testsBroken SQL, schema mismatches, non-idempotent loads

Deployment Strategy

Once tests pass, you need a strategy for getting DAGs to production. The two common approaches:

Git-sync: A sidecar container pulls from your main branch every 60 seconds. Merging to main is your deployment.

Image bake: DAGs are bundled into the Airflow Docker image at build time. Deploying means rolling out a new image version.

In either case, the CI pipeline is your safety net. No code reaches the branch that triggers deployment without passing all tests.

A Practical Testing Checklist

Not every DAG needs every type of test. Here is a prioritized approach:

For every Airflow project (non-negotiable):

  • DAG import test — catches 60% of production incidents
  • Linting in CI — catches typos and style issues

For critical pipelines (strongly recommended):

  • Structure tests for task existence and dependencies
  • Unit tests for transformation and business logic
  • Owner and tag validation

For complex pipelines (when the stakes are high):

  • Integration tests against a test database
  • Mock tests for external service interactions
  • Idempotency verification tests

Start with the import test. It takes 15 minutes to set up and catches a disproportionate number of problems. Add more tests as your DAGs grow in complexity and business importance.

Next Steps

Testing Airflow DAGs is not about achieving 100% coverage — it is about building confidence that your pipelines produce correct data. Start with validation tests, extract business logic into testable functions, and automate everything in CI.

  • Dynamic DAGs — Dynamic DAGs amplify the importance of testing, since one factory function controls many pipelines.
  • Best Practices — Testing is one pillar of production readiness; review the full checklist.
  • Production Deployment — Connect your CI pipeline to a proper deployment strategy.