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.