Testing

Airflow TaskFlow API Testing: pytest Patterns for Modern DAGs

HelpMeTest

18 Jul 2026 — 4 min read

The TaskFlow API (introduced in Airflow 2.0) transforms DAG authoring: instead of operator-heavy boilerplate, you write plain Python functions decorated with @task. This makes Airflow code significantly more testable — task logic is just a function, and functions can be unit tested.

This post covers testing patterns specific to the TaskFlow API: testing task functions directly, mocking XCom, testing dynamic task mapping, and validating DAG structure in CI.

TaskFlow Basics and What's Testable

# dags/order_processing_dag.py
from airflow.decorators import dag, task
from datetime import datetime

@dag(schedule="@daily", start_date=datetime(2026, 1, 1))
def order_processing():
    
    @task
    def extract_orders(execution_date: str) -> list[dict]:
        """Pull orders from source API."""
        # real implementation calls API
        return fetch_orders_from_api(execution_date)
    
    @task
    def validate_orders(orders: list[dict]) -> list[dict]:
        """Filter out invalid orders."""
        return [o for o in orders if o.get("amount", 0) > 0 and o.get("customer_id")]
    
    @task
    def compute_revenue(orders: list[dict]) -> dict:
        """Aggregate revenue by customer."""
        result = {}
        for order in orders:
            cid = order["customer_id"]
            result[cid] = result.get(cid, 0) + order["amount"]
        return result
    
    @task
    def load_to_warehouse(revenue: dict) -> None:
        write_to_bigquery(revenue)
    
    raw = extract_orders("{{ ds }}")
    valid = validate_orders(raw)
    rev = compute_revenue(valid)
    load_to_warehouse(rev)

dag = order_processing()

With TaskFlow, validate_orders and compute_revenue are testable as plain Python functions — no Airflow needed.

Testing Task Functions Directly

# tests/test_order_processing_tasks.py
import pytest
from dags.order_processing_dag import order_processing

# Get task functions from the DAG
dag_obj = order_processing()

def test_validate_orders_removes_zero_amount():
    validate = dag_obj.task_dict["validate_orders"].python_callable
    
    orders = [
        {"order_id": 1, "customer_id": "c1", "amount": 100.0},
        {"order_id": 2, "customer_id": "c2", "amount": 0.0},    # invalid
        {"order_id": 3, "customer_id": "c3", "amount": -5.0},   # invalid
        {"order_id": 4, "customer_id": "c4", "amount": 0.01},   # valid
    ]
    
    result = validate(orders)
    
    assert len(result) == 2
    assert all(o["amount"] > 0 for o in result)

def test_validate_orders_removes_missing_customer():
    validate = dag_obj.task_dict["validate_orders"].python_callable
    
    orders = [
        {"order_id": 1, "customer_id": "c1", "amount": 50.0},
        {"order_id": 2, "amount": 75.0},  # missing customer_id
        {"order_id": 3, "customer_id": None, "amount": 30.0},  # null customer_id
    ]
    
    result = validate(orders)
    
    assert len(result) == 1
    assert result[0]["order_id"] == 1

def test_compute_revenue_aggregates_per_customer():
    compute = dag_obj.task_dict["compute_revenue"].python_callable
    
    orders = [
        {"customer_id": "c1", "amount": 100.0},
        {"customer_id": "c1", "amount": 50.0},
        {"customer_id": "c2", "amount": 200.0},
    ]
    
    result = compute(orders)
    
    assert result["c1"] == 150.0
    assert result["c2"] == 200.0

def test_compute_revenue_handles_empty_list():
    compute = dag_obj.task_dict["compute_revenue"].python_callable
    
    result = compute([])
    
    assert result == {}

Using the TaskFlow Decorator Pattern

A cleaner approach extracts task functions before decorating:

# dags/order_processing_dag.py

# ✅ Extract business logic as pure functions
def _validate_orders(orders: list[dict]) -> list[dict]:
    return [o for o in orders if o.get("amount", 0) > 0 and o.get("customer_id")]

def _compute_revenue(orders: list[dict]) -> dict:
    result = {}
    for order in orders:
        cid = order["customer_id"]
        result[cid] = result.get(cid, 0) + order["amount"]
    return result

@dag(schedule="@daily", start_date=datetime(2026, 1, 1))
def order_processing():
    
    @task
    def validate_orders(orders):
        return _validate_orders(orders)  # delegates to testable function
    
    @task
    def compute_revenue(orders):
        return _compute_revenue(orders)
    # ...

# tests/test_order_processing.py — now trivially simple
from dags.order_processing_dag import _validate_orders, _compute_revenue

def test_validate_orders_removes_zero_amount():
    result = _validate_orders([
        {"order_id": 1, "customer_id": "c1", "amount": 100.0},
        {"order_id": 2, "customer_id": "c2", "amount": 0.0},
    ])
    assert len(result) == 1

def test_compute_revenue_aggregates_correctly():
    result = _compute_revenue([
        {"customer_id": "c1", "amount": 100.0},
        {"customer_id": "c1", "amount": 50.0},
    ])
    assert result["c1"] == 150.0

Testing with the Airflow Test Harness

For integration tests that need real Airflow infrastructure:

# tests/test_dag_integration.py
import pytest
from airflow.models import DagBag
from airflow.utils.state import State
from airflow.utils.types import DagRunType

@pytest.fixture
def dagbag():
    return DagBag(dag_folder="dags/", include_examples=False)

def test_dag_loads_without_errors(dagbag):
    assert len(dagbag.import_errors) == 0, \
        f"DAG import errors: {dagbag.import_errors}"

def test_order_processing_dag_exists(dagbag):
    dag = dagbag.get_dag("order_processing")
    assert dag is not None

def test_dag_task_count(dagbag):
    dag = dagbag.get_dag("order_processing")
    assert len(dag.tasks) == 4  # extract, validate, compute, load

def test_dag_has_correct_schedule(dagbag):
    dag = dagbag.get_dag("order_processing")
    assert dag.schedule_interval == "@daily"

Testing Dynamic Task Mapping

Dynamic task mapping (Airflow 2.3+) creates tasks at runtime based on data:

# dags/parallel_processing_dag.py
from airflow.decorators import dag, task

@dag(schedule="@daily", start_date=datetime(2026, 1, 1))
def parallel_processing():
    
    @task
    def get_partitions() -> list[str]:
        return ["2026-01-01", "2026-01-02", "2026-01-03"]
    
    @task
    def process_partition(partition: str) -> dict:
        """Processes one partition — runs in parallel for each partition."""
        return {"partition": partition, "rows_processed": load_partition(partition)}
    
    @task
    def aggregate_results(results: list[dict]) -> dict:
        total = sum(r["rows_processed"] for r in results)
        return {"total_rows": total, "partitions": len(results)}
    
    partitions = get_partitions()
    processed = process_partition.expand(partition=partitions)
    aggregate_results(processed)

dag = parallel_processing()

# tests/test_parallel_processing.py
from dags.parallel_processing_dag import parallel_processing

def test_process_partition_extracts_correct_data():
    # Test the individual partition processor — the mapped task function
    dag_obj = parallel_processing()
    process = dag_obj.task_dict["process_partition"].python_callable
    
    # Mock the underlying data load
    with patch("dags.parallel_processing_dag.load_partition") as mock_load:
        mock_load.return_value = 1000
        
        result = process("2026-01-01")
        
        assert result["partition"] == "2026-01-01"
        assert result["rows_processed"] == 1000
        mock_load.assert_called_once_with("2026-01-01")

def test_aggregate_results_sums_correctly():
    dag_obj = parallel_processing()
    aggregate = dag_obj.task_dict["aggregate_results"].python_callable
    
    results = [
        {"partition": "2026-01-01", "rows_processed": 1000},
        {"partition": "2026-01-02", "rows_processed": 1500},
        {"partition": "2026-01-03", "rows_processed": 2000},
    ]
    
    result = aggregate(results)
    
    assert result["total_rows"] == 4500
    assert result["partitions"] == 3

Testing Sensors

Sensors wait for external conditions. Test them with mocked conditions:

# tests/test_sensors.py
from unittest.mock import patch, MagicMock
from airflow.sensors.base import BaseSensorOperator

class TestS3FileSensor:
    
    def test_poke_returns_true_when_file_exists(self):
        sensor = S3KeySensor(
            task_id="wait_for_file",
            bucket_name="my-bucket",
            bucket_key="data/orders_2026-01-01.parquet",
            dag=dag
        )
        
        with patch("boto3.client") as mock_boto:
            mock_s3 = MagicMock()
            mock_boto.return_value = mock_s3
            mock_s3.head_object.return_value = {"ContentLength": 1024}
            
            result = sensor.poke(context={})
            
            assert result is True
    
    def test_poke_returns_false_when_file_missing(self):
        sensor = S3KeySensor(
            task_id="wait_for_file",
            bucket_name="my-bucket",
            bucket_key="data/missing.parquet",
            dag=dag
        )
        
        with patch("boto3.client") as mock_boto:
            mock_s3 = MagicMock()
            mock_boto.return_value = mock_s3
            from botocore.exceptions import ClientError
            mock_s3.head_object.side_effect = ClientError(
                {"Error": {"Code": "404"}}, "HeadObject")
            
            result = sensor.poke(context={})
            
            assert result is False

CI Configuration

# .github/workflows/airflow-tests.yml
name: Airflow Tests
on: [push, pull_request]

jobs:
  test:
    runs-on: ubuntu-latest
    services:
      postgres:
        image: postgres:15
        env:
          POSTGRES_USER: airflow
          POSTGRES_PASSWORD: airflow
          POSTGRES_DB: airflow
        ports:
          - 5432:5432
    
    env:
      AIRFLOW__DATABASE__SQL_ALCHEMY_CONN: postgresql+psycopg2://airflow:airflow@localhost/airflow
      AIRFLOW__CORE__EXECUTOR: LocalExecutor
      AIRFLOW__CORE__LOAD_EXAMPLES: "false"
    
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: '3.11'
      
      - name: Install Airflow and dependencies
        run: |
          pip install apache-airflow==2.9.0 pytest pytest-cov
          airflow db init
      
      - name: Run unit tests (no Airflow DB needed)
        run: pytest tests/unit/ -v
      
      - name: Run integration tests (needs Airflow DB)
        run: pytest tests/integration/ -v

The unit tests (testing pure task functions) run in seconds without any infrastructure. The integration tests (testing DAG loading and structure) need the Airflow metadata DB but still don't need a running Airflow environment.

The pattern that makes this work: extract all business logic into pure functions, wrap them in @task decorators that only handle Airflow plumbing. The business logic tests are fast. The DAG structure tests catch import errors and broken dependencies before they reach production.