Unit testing DAGs in Apache Airflow - Commands & Configuration

Introduction

Unit testing DAGs helps you check if your workflow definitions in Airflow are correct before running them. It finds mistakes early so your scheduled tasks run smoothly without errors.

When you want to verify that your DAG structure is valid and all tasks are connected properly.

When you add new tasks or change dependencies and want to ensure nothing breaks.

When you want to check that task parameters like retries and schedules are set as expected.

When you want to catch syntax errors or missing imports in your DAG files before deployment.

When you want to automate testing of your workflows as part of your development process.

Config File - test_example_dag.py

test_example_dag.py

import unittest
from airflow.models import DagBag

class TestExampleDAG(unittest.TestCase):
    def setUp(self):
        self.dagbag = DagBag(dag_folder="./dags", include_examples=False)

    def test_no_import_errors(self):
        self.assertEqual(len(self.dagbag.import_errors), 0, f"Import errors found: {self.dagbag.import_errors}")

    def test_dag_loaded(self):
        dag = self.dagbag.get_dag(dag_id="example_dag")
        self.assertIsNotNone(dag)
        self.assertEqual(dag.dag_id, "example_dag")

    def test_task_count(self):
        dag = self.dagbag.get_dag(dag_id="example_dag")
        self.assertEqual(len(dag.tasks), 3)

if __name__ == "__main__":
    unittest.main()

This Python test file uses unittest to load DAGs from the ./dags folder.

setUp: Prepares the DAG bag without example DAGs.

test_no_import_errors: Checks that all DAG files load without errors.

test_dag_loaded: Confirms the DAG with ID example_dag exists.

test_task_count: Verifies the DAG has exactly 3 tasks.

Commands

Runs the unit tests to check if the DAGs load correctly and have the expected structure.

Terminal

python test_example_dag.py

Expected Output

...

----------------------------------------------------------------------

Ran 3 tests in 0.123s



OK

Lists all DAGs currently available in Airflow to verify the DAG is recognized by the system.

Terminal

airflow dags list

Expected Output

example_dag

Key Concept

If you remember nothing else from unit testing DAGs, remember: always check for import errors and confirm your DAG loads with the expected tasks before running it.

Common Mistakes

Not checking for import errors in DAG files.

Import errors cause DAGs to not load, which breaks scheduling and execution.

Use DagBag's import_errors attribute in tests to catch and fix import problems early.

Assuming the DAG exists without verifying it in tests.

If the DAG ID is wrong or the file is missing, tests will pass incorrectly or fail later in production.

Explicitly get the DAG by ID in tests and assert it is not None.

Not verifying the number or names of tasks in the DAG.

Missing or extra tasks can cause workflow failures or unexpected behavior.

Write tests that check the exact count and identity of tasks in the DAG.

Summary

Use DagBag in unit tests to load DAGs and check for import errors.

Write tests to confirm the DAG exists and has the expected tasks.

Run tests with Python to catch issues before deploying workflows.

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