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

Need a 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

Need a 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

Need a 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

Need a hint?

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