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.

Practice

1. What is the main purpose of unit testing DAGs in Airflow?

easy

A. To catch errors early before running workflows

B. To deploy DAGs to production automatically

C. To monitor task execution times

D. To schedule DAG runs at specific times

Solution

1. Step 1: Understand unit testing purpose

   Unit testing is done to find bugs or errors early in code before actual execution.

2. Step 2: Apply this to Airflow DAGs

   Unit testing DAGs helps catch errors in DAG definitions before running workflows in Airflow.

3. Final Answer:

   To catch errors early before running workflows -> Option A

4. Quick Check:

   Unit testing = catch errors early [OK]

Hint: Unit tests find bugs before running DAGs [OK]

Common Mistakes:

• Confusing unit testing with deployment
• Thinking unit tests schedule DAGs
• Mixing monitoring with testing

2. Which Python code snippet correctly loads DAGs for unit testing using Airflow's DagBag?

easy

A. dagbag = DagBag.load('/path/to/dags')

B. dagbag = DagBag(dag_folder='/path/to/dags')

C. dagbag = DagBag(folder='/path/to/dags')

D. dagbag = DagBag.get('/path/to/dags')

Solution

1. Step 1: Recall DagBag constructor usage

   DagBag takes a keyword argument dag_folder to specify DAGs folder path.

2. Step 2: Check each option

   Only dagbag = DagBag(dag_folder='/path/to/dags') uses correct syntax. Others use invalid methods (A,D) or incorrect keyword argument like folder (B).

3. Final Answer:

   dagbag = DagBag(dag_folder='/path/to/dags') -> Option B

4. Quick Check:

   DagBag(dag_folder=...) correct syntax [OK]

Hint: Use dag_folder keyword to load DAGs in DagBag [OK]

Common Mistakes:

• Using incorrect keyword argument (e.g., folder)
• Using non-existent methods like load() or get()
• Omitting dag_folder keyword

3. Given this test code snippet, what will be the output if the DAG named 'example_dag' has 3 tasks?

dagbag = DagBag(dag_folder='dags')
dag = dagbag.get_dag('example_dag')
print(len(dag.tasks))

medium

A. None

B. 0

C. 3

D. Error: DAG not found

Solution

1. Step 1: Load DAG using DagBag

   DagBag loads DAGs from 'dags' folder; get_dag('example_dag') returns the DAG object if found.

2. Step 2: Count tasks in the DAG

   len(dag.tasks) returns the number of tasks in the DAG, which is 3 as given.

3. Final Answer:

   3 -> Option C

4. Quick Check:

   len(dag.tasks) = 3 [OK]

Hint: len(dag.tasks) gives task count [OK]

Common Mistakes:

• Assuming get_dag returns None if DAG exists
• Confusing tasks with task instances
• Expecting error if DAG is present

4. You wrote this test code but it raises an AttributeError: 'NoneType' object has no attribute 'tasks'. What is the likely cause?

dagbag = DagBag(dag_folder='dags')
dag = dagbag.get_dag('missing_dag')
print(len(dag.tasks))

medium

A. The DAG 'missing_dag' does not exist in the folder

B. DagBag cannot load DAGs from 'dags' folder

C. len() cannot be used on dag.tasks

D. get_dag() returns a list, not a DAG object

Solution

1. Step 1: Analyze get_dag return value

   get_dag returns None if the DAG name is not found in the loaded DAGs.

2. Step 2: Understand AttributeError cause

   Trying to access dag.tasks when dag is None causes AttributeError.

3. Final Answer:

   The DAG 'missing_dag' does not exist in the folder -> Option A

4. Quick Check:

   get_dag returns None if DAG missing [OK]

Hint: Check if get_dag returns None before accessing tasks [OK]

Common Mistakes:

• Assuming get_dag always returns a DAG
• Thinking len() is invalid on tasks list
• Believing DagBag can't load DAGs from folder

5. You want to write a unit test to check that your DAG named 'data_pipeline' has exactly 5 tasks and no import errors. Which code snippet correctly performs this test?

hard

A. dagbag = DagBag(dag_folder='dags')
assert dagbag.import_errors == {}
dag = dagbag.get_dag('data_pipeline')
assert len(dag.tasks) == 0

B. dagbag = DagBag('dags')
assert dagbag.import_errors == None
dag = dagbag.get_dag('data_pipeline')
assert len(dag.tasks) == 5

C. dagbag = DagBag(dag_folder='dags')
assert dagbag.import_errors != {}
dag = dagbag.get_dag('data_pipeline')
assert len(dag.tasks) == 5

D. dagbag = DagBag(dag_folder='dags')
assert dagbag.import_errors == {}
dag = dagbag.get_dag('data_pipeline')
assert len(dag.tasks) == 5

Solution

1. Step 1: Check import errors correctly

   dagbag.import_errors is a dict; empty dict {} means no import errors.

2. Step 2: Verify task count assertion

   len(dag.tasks) == 5 checks that the DAG has exactly 5 tasks.

3. Final Answer:

   dagbag = DagBag(dag_folder='dags') assert dagbag.import_errors == {} dag = dagbag.get_dag('data_pipeline') assert len(dag.tasks) == 5 -> Option D

4. Quick Check:

   import_errors == {} and task count 5 [OK]

Hint: Assert import_errors == {} and task count matches [OK]

Common Mistakes:

• Checking import_errors against None instead of {}
• Asserting wrong task count
• Confusing import_errors presence logic