Testing custom operators in Apache Airflow - Mini Project: Build & Apply

Testing Custom Operators in Airflow

📖 Scenario: You are working with Apache Airflow to automate workflows. You have created a custom operator to perform a simple task. Now, you want to write a test to ensure your custom operator works correctly before using it in production.

🎯 Goal: Build a simple test for a custom Airflow operator that prints a message. You will create the operator, set up a test DAG, and write a test function to verify the operator runs without errors.

📋 What You'll Learn

Create a custom operator class called PrintOperator that inherits from BaseOperator.

Add an execute method to PrintOperator that prints the message 'Hello from PrintOperator'.

Create a DAG named test_print_operator_dag with a single task using PrintOperator.

Write a test function called test_print_operator_execute that runs the operator's execute method and checks it completes without error.

💡 Why This Matters

🌍 Real World

Custom operators let you extend Airflow to automate tasks specific to your business or environment, such as data processing, notifications, or system commands.

💼 Career

Knowing how to create and test custom operators is important for Airflow developers and DevOps engineers to build reliable, maintainable workflows.

1

Create the custom operator class

Create a class called PrintOperator that inherits from BaseOperator. Inside it, define an execute method that prints the exact message 'Hello from PrintOperator'.

Apache Airflow

from airflow.models import BaseOperator

# Define PrintOperator class here
# Your code here

Hint

Remember to inherit from BaseOperator and define execute(self, context).

2

Create a DAG with the custom operator task

Create a DAG named test_print_operator_dag using DAG from airflow. Inside the DAG context, create a task called print_task using PrintOperator with task_id='print_task'.

Apache Airflow

from airflow.models import BaseOperator

class PrintOperator(BaseOperator):
    def execute(self, context):
        print('Hello from PrintOperator')

from airflow import DAG
from datetime import datetime

# Create DAG and task here
# Your code here

Hint

Use with DAG(...) to create the DAG and instantiate PrintOperator inside it.

3

Write a test function to run the operator

Write a function called test_print_operator_execute that creates an instance of PrintOperator with task_id='test_task' and calls its execute method with an empty dictionary as the context.

Apache Airflow

from airflow.models import BaseOperator

class PrintOperator(BaseOperator):
    def execute(self, context):
        print('Hello from PrintOperator')

from airflow import DAG
from datetime import datetime

with DAG('test_print_operator_dag', start_date=datetime(2024, 1, 1), schedule_interval=None) as dag:
    print_task = PrintOperator(task_id='print_task')

# Define test function here
# Your code here

Hint

The test function should create the operator and call execute with an empty context.

4

Run the test function and print confirmation

Call the test_print_operator_execute function and then print the exact message 'Test completed successfully'.

Apache Airflow

from airflow.models import BaseOperator

class PrintOperator(BaseOperator):
    def execute(self, context):
        print('Hello from PrintOperator')

from airflow import DAG
from datetime import datetime

with DAG('test_print_operator_dag', start_date=datetime(2024, 1, 1), schedule_interval=None) as dag:
    print_task = PrintOperator(task_id='print_task')

def test_print_operator_execute():
    operator = PrintOperator(task_id='test_task')
    operator.execute(context={})

# Run the test function and print confirmation
# Your code here

Hint

Call the test function and then print the confirmation message exactly as shown.

Practice

1. What is the main reason to write tests for custom Airflow operators?

easy

A. To ensure the operator works as expected before running in production

B. To make the DAG run faster

C. To reduce the number of tasks in the DAG

D. To automatically generate documentation

Solution

1. Step 1: Understand the purpose of testing

   Testing is done to verify that code behaves correctly and avoids errors in production.

2. Step 2: Apply this to custom operators

   Custom operators are pieces of code that need validation to ensure they perform their tasks as intended.

3. Final Answer:

   To ensure the operator works as expected before running in production -> Option A

4. Quick Check:

   Testing = Verify correctness [OK]

Hint: Testing checks correctness before production use [OK]

Common Mistakes:

• Thinking tests speed up DAG execution
• Confusing testing with documentation generation
• Believing tests reduce task count

2. Which Python module is commonly used to write tests for Airflow custom operators?

easy

A. requests

B. flask

C. unittest

D. pandas

Solution

1. Step 1: Identify common testing modules in Python

   Python's built-in module for testing is unittest, widely used for writing test cases.

2. Step 2: Match with Airflow testing practices

   Airflow tests often use unittest to create test classes and methods for operators.

3. Final Answer:

   unittest -> Option C

4. Quick Check:

   Python testing = unittest [OK]

Hint: Use unittest for Python testing by default [OK]

Common Mistakes:

• Choosing unrelated modules like flask or pandas
• Confusing HTTP libraries with testing tools

3. Given this test snippet for a custom operator:

from airflow.models import DagBag

def test_operator_load():
    dagbag = DagBag()
    assert 'my_custom_dag' in dagbag.dags

What does this test check?

medium

A. If the DAG named 'my_custom_dag' is loaded successfully

B. If the operator runs without errors

C. If the DAG tasks execute in order

D. If the DAG has no syntax errors

Solution

1. Step 1: Understand DagBag usage

   DagBag loads all DAGs from the Airflow folder and stores them in dagbag.dags dictionary.

2. Step 2: Analyze the assertion

   The test asserts that 'my_custom_dag' is a key in dagbag.dags, meaning it was loaded.

3. Final Answer:

   If the DAG named 'my_custom_dag' is loaded successfully -> Option A

4. Quick Check:

   DagBag loads DAGs = check DAG presence [OK]

Hint: DagBag.dags keys are loaded DAG ids [OK]

Common Mistakes:

• Thinking it tests operator execution
• Confusing DAG loading with task order
• Assuming it checks syntax errors directly

4. You wrote a test for a custom operator but it fails with AttributeError: 'NoneType' object has no attribute 'execute'. What is the likely cause?

medium

A. The DAG file has syntax errors

B. The operator instance was not created properly in the test

C. The Airflow scheduler is not running

D. The test file is missing import statements

Solution

1. Step 1: Understand the error meaning

   AttributeError on 'NoneType' means the variable expected to be an object is actually None.

2. Step 2: Relate to operator instance creation

   If the operator instance was not created or assigned, calling execute() on None causes this error.

3. Final Answer:

   The operator instance was not created properly in the test -> Option B

4. Quick Check:

   None.execute error = missing instance [OK]

Hint: Check if operator object is None before calling methods [OK]

Common Mistakes:

• Blaming DAG syntax errors for this runtime error
• Thinking scheduler status affects unit tests
• Ignoring missing or wrong object creation

5. You want to test a custom operator that makes an HTTP request. Which approach is best to avoid real network calls during tests?

hard

A. Write the test to call the real HTTP endpoint every time

B. Run the test only when connected to the internet

C. Skip the test if the HTTP server is down

D. Use mocking to replace the HTTP call with a fake response

Solution

1. Step 1: Understand testing external calls

   Real network calls can be slow, flaky, or cause side effects, so they should be avoided in unit tests.

2. Step 2: Use mocking for external dependencies

   Mocking replaces the real HTTP call with a controlled fake response, making tests fast and reliable.

3. Final Answer:

   Use mocking to replace the HTTP call with a fake response -> Option D

4. Quick Check:

   Mock external calls = reliable tests [OK]

Hint: Mock external calls to avoid real network in tests [OK]

Common Mistakes:

• Running tests dependent on internet availability
• Skipping tests instead of fixing design
• Calling real endpoints causing flaky tests