Unit testing DAGs in Apache Airflow - Step-by-Step Execution

Process Flow - Unit testing DAGs

Write DAG code

↓

Write unit test code

↓

Run unit tests

↓

Check test results

↓

Deploy

↩Back to Write unit test code

This flow shows how you write DAGs, create unit tests for them, run tests, and then either deploy or fix issues based on results.

Execution Sample

Apache Airflow

def test_dag_structure():
    dag = get_dag('example_dag')
    assert dag is not None
    assert len(dag.tasks) == 3
    assert dag.tasks[0].task_id == 'start_task'

This test checks that the DAG named 'example_dag' exists, has 3 tasks, and the first task is 'start_task'.

Process Table

Step	Action	Evaluation	Result
1	Call get_dag('example_dag')	DAG object found?	True
2	Check dag is not None	dag is not None	Pass
3	Check number of tasks	len(dag.tasks) == 3	Pass
4	Check first task id	dag.tasks[0].task_id == 'start_task'	Pass
5	All assertions passed	Test result	Pass

💡 All assertions passed, test completes successfully

Status Tracker

Variable	Start	After Step 1	After Step 2	After Step 3	After Step 4	Final
dag	None	DAG object for 'example_dag'	Same	Same	Same	Same
len(dag.tasks)	N/A	N/A	N/A	3	Same	Same
dag.tasks[0].task_id	N/A	N/A	N/A	N/A	'start_task'	'start_task'

Key Moments - 3 Insights

Why do we check if dag is not None before other assertions?

Because if dag is None, accessing dag.tasks would cause an error. Step 2 in the execution_table ensures the DAG exists before checking tasks.

What happens if the number of tasks is not 3?

The assertion in Step 3 would fail, stopping the test early and indicating the DAG structure is incorrect.

Why check the first task's task_id?

To verify the DAG tasks are correctly named and ordered, ensuring the workflow is as expected (Step 4).

Visual Quiz - 3 Questions

Look at the execution_table, what is the value of 'dag' after Step 1?

A None

B 3

C DAG object for 'example_dag'

D 'start_task'

💡 Hint

Check the variable_tracker row for 'dag' at 'After Step 1'

At which step does the test verify the number of tasks in the DAG?

A Step 3

B Step 4

C Step 2

D Step 5

💡 Hint

Look at the 'Action' column in execution_table for checking task count

If the DAG had 4 tasks instead of 3, what would happen in the execution_table?

A Step 4 would fail

B Step 3 would fail

C Step 3 would pass

D Step 5 would pass

💡 Hint

Step 3 checks if len(dag.tasks) == 3; if not, assertion fails there

Concept Snapshot

Unit testing DAGs in Airflow:
- Write tests to check DAG existence and structure
- Use assertions to verify tasks count and IDs
- Run tests before deployment
- Fix errors if tests fail
- Ensures DAG correctness early

Full Transcript

Unit testing DAGs means writing small test functions that check if your DAG exists and has the right tasks. First, you get the DAG object by name. Then you check if it is not None to avoid errors. Next, you verify the number of tasks matches what you expect. Finally, you check specific task IDs to confirm the workflow order. If all assertions pass, the test succeeds and you can deploy safely. If any assertion fails, you fix the DAG or tests before moving on. This process helps catch mistakes early and keeps your workflows reliable.