Testing custom operators in Apache Airflow - Step-by-Step Execution

Process Flow - Testing custom operators

Write custom operator code

↓

Write test case for operator

↓

Run test with test framework

↓

Test framework executes operator code

↓

Check operator output and side effects

↓

Pass or Fail test

↓

Fix bugs if Fail, else Done

This flow shows writing a custom operator, creating tests, running them, and checking results to ensure the operator works correctly.

Execution Sample

Apache Airflow

from airflow.models import BaseOperator

class MyOperator(BaseOperator):
    def execute(self, context):
        return 'done'

# Test
op = MyOperator(task_id='test')
assert op.execute({}) == 'done'

This code defines a simple custom operator and tests its execute method returns 'done'.

Process Table

Step	Action	Input	Output	Result
1	Create MyOperator instance	task_id='test'	MyOperator object	Instance created
2	Call execute method	context={}	return value	'done' returned
3	Assert output	'done' == 'done'	True	Test passes

💡 Test passes because execute returns expected 'done'

Status Tracker

Variable	Start	After Step 1	After Step 2	Final
op	None	MyOperator instance	MyOperator instance	MyOperator instance
result	None	None	'done'	'done'

Key Moments - 3 Insights

Why do we create an instance of the operator before testing?

Because the execute method is an instance method, so we must create the operator object first (see Step 1 in execution_table).

What does the assert statement check in the test?

It checks that the execute method returns the expected output 'done' (see Step 3 in execution_table).

Why do we pass an empty dictionary as context?

Because execute expects a context dictionary; here we test basic behavior without extra data (Step 2 input).

Visual Quiz - 3 Questions

Look at the execution table, what is the output of the execute method at Step 2?

A MyOperator instance

B None

C 'done'

D Error

💡 Hint

Check the Output column at Step 2 in the execution_table

At which step does the test confirm the operator works as expected?

A Step 1

B Step 3

C Step 2

D After Step 3

💡 Hint

Look at the Result column where the test passes in execution_table

If the execute method returned 'fail' instead of 'done', what would change in the execution table?

A Step 3 Result would be 'Test fails'

B Step 3 Result would be 'Test passes'

C Step 2 Output would be 'done'

D Step 1 Action would change

💡 Hint

Consider the assert check result in Step 3 of execution_table

Concept Snapshot

Testing custom operators in Airflow:
- Define operator class with execute method
- Create operator instance in test
- Call execute with context dict
- Assert execute returns expected result
- Fix code if test fails

Full Transcript

To test a custom Airflow operator, first write the operator class with an execute method. Then write a test that creates an instance of this operator. Call the execute method with a context dictionary. Check the output matches the expected result using an assert statement. If the test passes, the operator works as intended. If it fails, fix the operator code and test again.

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