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.