Testing custom operators in Apache Airflow - Time & Space Complexity
When testing custom operators in Airflow, it's important to understand how the test execution time grows as the number of test cases or operator complexity increases.
We want to know how the time to run tests changes when we add more tests or more complex logic.
Analyze the time complexity of the following test code snippet.
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
class MyCustomOperator(BaseOperator):
@apply_defaults
def __init__(self, param, **kwargs):
super().__init__(**kwargs)
self.param = param
def execute(self, context):
return self.param * 2
def test_my_custom_operator():
operator = MyCustomOperator(task_id='test', param=5)
result = operator.execute(context={})
assert result == 10
This code defines a simple custom operator and a test function that runs the operator's execute method once and checks the result.
Look for loops or repeated calls in the test.
- Primary operation: Single call to
executemethod of the operator. - How many times: Exactly once per test run.
Since the test runs the operator once, the time grows directly with the number of tests.
| Input Size (number of tests) | Approx. Operations |
|---|---|
| 10 | 10 calls to execute |
| 100 | 100 calls to execute |
| 1000 | 1000 calls to execute |
Pattern observation: The total test time grows linearly as you add more tests.
Time Complexity: O(n)
This means the test time increases in direct proportion to the number of tests you run.
[X] Wrong: "Running more tests won't affect total test time much because each test is small."
[OK] Correct: Even small tests add up, so doubling tests roughly doubles total time.
Understanding how test time grows helps you write efficient tests and manage test suites well in real projects.
"What if we added a loop inside the execute method that runs proportional to an input size? How would the time complexity change?"
Practice
Solution
Step 1: Understand the purpose of testing
Testing is done to verify that code behaves correctly and avoids errors in production.Step 2: Apply this to custom operators
Custom operators are pieces of code that need validation to ensure they perform their tasks as intended.Final Answer:
To ensure the operator works as expected before running in production -> Option AQuick Check:
Testing = Verify correctness [OK]
- Thinking tests speed up DAG execution
- Confusing testing with documentation generation
- Believing tests reduce task count
Solution
Step 1: Identify common testing modules in Python
Python's built-in module for testing is unittest, widely used for writing test cases.Step 2: Match with Airflow testing practices
Airflow tests often use unittest to create test classes and methods for operators.Final Answer:
unittest -> Option CQuick Check:
Python testing = unittest [OK]
- Choosing unrelated modules like flask or pandas
- Confusing HTTP libraries with testing tools
from airflow.models import DagBag
def test_operator_load():
dagbag = DagBag()
assert 'my_custom_dag' in dagbag.dagsWhat does this test check?
Solution
Step 1: Understand DagBag usage
DagBag loads all DAGs from the Airflow folder and stores them in dagbag.dags dictionary.Step 2: Analyze the assertion
The test asserts that 'my_custom_dag' is a key in dagbag.dags, meaning it was loaded.Final Answer:
If the DAG named 'my_custom_dag' is loaded successfully -> Option AQuick Check:
DagBag loads DAGs = check DAG presence [OK]
- Thinking it tests operator execution
- Confusing DAG loading with task order
- Assuming it checks syntax errors directly
Solution
Step 1: Understand the error meaning
AttributeError on 'NoneType' means the variable expected to be an object is actually None.Step 2: Relate to operator instance creation
If the operator instance was not created or assigned, calling execute() on None causes this error.Final Answer:
The operator instance was not created properly in the test -> Option BQuick Check:
None.execute error = missing instance [OK]
- Blaming DAG syntax errors for this runtime error
- Thinking scheduler status affects unit tests
- Ignoring missing or wrong object creation
Solution
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.Step 2: Use mocking for external dependencies
Mocking replaces the real HTTP call with a controlled fake response, making tests fast and reliable.Final Answer:
Use mocking to replace the HTTP call with a fake response -> Option DQuick Check:
Mock external calls = reliable tests [OK]
- Running tests dependent on internet availability
- Skipping tests instead of fixing design
- Calling real endpoints causing flaky tests
Ready to prove it?
Pass the test and this pattern moves unknown
