Unit testing DAGs in Apache Airflow - Mini Project: Build & Apply

Unit Testing Airflow DAGs

📖 Scenario: You work as a data engineer. You create workflows called DAGs in Apache Airflow to automate tasks. To keep your workflows reliable, you want to write simple tests that check if your DAGs are set up correctly.

🎯 Goal: Build a small unit test for an Airflow DAG. You will create a DAG, add a configuration variable, write a test function to check the DAG's tasks, and finally run the test to see the result.

📋 What You'll Learn

Create a DAG object with a specific DAG ID

Add a default argument dictionary with a start date

Write a test function that checks the DAG ID and task count

Print the test result output

💡 Why This Matters

🌍 Real World

Data engineers use Airflow DAGs to automate data workflows. Unit testing DAGs helps catch errors early and keeps pipelines reliable.

💼 Career

Knowing how to write unit tests for Airflow DAGs is a valuable skill for data engineers and DevOps professionals working with data pipelines.

1

Create a simple Airflow DAG

Import DAG from airflow.models and create a DAG object called test_dag with the DAG ID 'example_dag' and an empty default_args dictionary.

Apache Airflow

# Import DAG and create test_dag with dag_id 'example_dag' and empty default_args
# Your code here

Hint

Use DAG(dag_id='example_dag', default_args=default_args) to create the DAG.

2

Add a start date to default_args

Add a start_date key to the default_args dictionary with the value datetime(2024, 1, 1). Import datetime from the datetime module.

Apache Airflow

from airflow.models import DAG
from datetime import datetime

default_args = {}
# Add start_date=datetime(2024, 1, 1) to default_args below
# Your code here

test_dag = DAG(dag_id='example_dag', default_args=default_args)

Hint

Use default_args = {'start_date': datetime(2024, 1, 1)}.

3

Write a test function to check the DAG

Define a function called test_dag_structure that checks if test_dag.dag_id equals 'example_dag' and if the number of tasks in test_dag.tasks is zero. Return True if both conditions are met, otherwise False.

Apache Airflow

from airflow.models import DAG
from datetime import datetime

default_args = {'start_date': datetime(2024, 1, 1)}
test_dag = DAG(dag_id='example_dag', default_args=default_args)

# Define test_dag_structure function below
# Your code here

Hint

Use def test_dag_structure(): and check conditions with if.

4

Run the test and print the result

Call the function test_dag_structure() and print its result using print().

Apache Airflow

from airflow.models import DAG
from datetime import datetime

default_args = {'start_date': datetime(2024, 1, 1)}
test_dag = DAG(dag_id='example_dag', default_args=default_args)

def test_dag_structure():
    if test_dag.dag_id == 'example_dag' and len(test_dag.tasks) == 0:
        return True
    else:
        return False

# Call test_dag_structure and print the result below
# Your code here

Hint

Use print(test_dag_structure()) to show the test result.

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