Testing custom operators in Apache Airflow - Commands & Configuration

Introduction

When you create your own custom operator in Airflow, you need to test it to make sure it works correctly before using it in real workflows. Testing helps catch mistakes early and ensures your operator behaves as expected.

When you build a new operator to perform a specific task in your data pipeline.

When you want to verify that your operator handles inputs and outputs correctly.

When you need to check that your operator interacts properly with other Airflow components.

When you want to prevent errors in production by catching bugs during development.

When you update an existing operator and want to confirm it still works as intended.

Config File - test_custom_operator.py

test_custom_operator.py

from airflow.models import DAG
from airflow.utils.dates import days_ago
from airflow.operators.python import PythonOperator
import unittest

# Example custom operator
class MyCustomOperator(PythonOperator):
    def execute(self, context):
        return 'custom operator executed'

class TestMyCustomOperator(unittest.TestCase):
    def setUp(self):
        self.dag = DAG(dag_id='test_dag', start_date=days_ago(1))

    def test_execute(self):
        operator = MyCustomOperator(
            task_id='test_task',
            python_callable=lambda: 'custom operator executed',
            dag=self.dag
        )
        result = operator.execute(context={})
        self.assertEqual(result, 'custom operator executed')

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

This Python file contains a simple custom operator class MyCustomOperator that extends Airflow's PythonOperator. It overrides the execute method to return a fixed string.

The TestMyCustomOperator class uses Python's unittest framework to test the operator. It creates a DAG for context, runs the operator's execute method, and checks that the output matches the expected string.

This test file can be run directly to verify the custom operator works as intended.

Commands

Run the test file to execute the unit test for the custom operator. This checks if the operator returns the expected result.

Terminal

python test_custom_operator.py

Expected Output

.

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

Ran 1 test in 0.001s



OK

List all DAGs in Airflow to confirm the test DAG is recognized by Airflow after adding the custom operator.

Terminal

airflow dags list

Expected Output

dag_id

-------

test_dag

Key Concept

If you remember nothing else from this pattern, remember: always write unit tests that run your custom operator's execute method and check its output to catch errors early.

Common Mistakes

Not providing a DAG context when creating the operator instance in tests.

Airflow operators require a DAG context; without it, the operator may fail or behave unexpectedly.

Always create a DAG object and pass it to the operator during testing.

Not running the execute method directly in tests and instead relying on Airflow scheduler.

The scheduler runs asynchronously and is harder to test; unit tests should run execute directly for fast feedback.

Call the execute method directly with a context dictionary in your unit tests.

Ignoring the output of the execute method in tests.

Without checking output, you cannot confirm the operator behaves as expected.

Use assertions to compare the execute method's return value to the expected result.

Summary

Create a test Python file that defines your custom operator and a unittest class.

Write a test method that creates the operator with a DAG context and calls execute.

Run the test file with Python to verify the operator works as expected.

Use airflow CLI commands to check your DAGs are loaded correctly.

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