Testing custom operators in Apache Airflow - Practice Problems & Coding Challenges

Challenge - 5 Problems

💻 Command Output

intermediate

Output of a simple custom operator execution

Given the following Airflow custom operator code snippet, what will be the output when the operator runs successfully?

Apache Airflow

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults

class HelloOperator(BaseOperator):
    @apply_defaults
    def __init__(self, name: str, **kwargs):
        super().__init__(**kwargs)
        self.name = name

    def execute(self, context):
        message = f"Hello, {self.name}!"
        print(message)
        return message

# Assume this operator is instantiated and run with name='Airflow'

A Hello, Airflow!

B Hello, name!

C Error: missing required argument

D None

💡 Hint

Look at the execute method and what it prints and returns.

The execute method prints and returns the message with the name passed during initialization. Since name='Airflow', the output is Hello, Airflow!.

❓ Configuration

intermediate

Correct DAG configuration to use a custom operator

Which DAG configuration snippet correctly uses a custom operator MyOperator with task_id task1?

Apache Airflow

from airflow import DAG
from datetime import datetime
from my_operators import MyOperator

default_args = {
    'start_date': datetime(2024, 1, 1),
}

dag = DAG('my_dag', default_args=default_args, schedule_interval='@daily')

# Fill in the correct task definition below

A task1 = MyOperator(task_id='task1')

B task1 = MyOperator(dag=dag)

C task1 = MyOperator('task1', dag=dag)

D task1 = MyOperator(task_id='task1', dag=dag)

💡 Hint

Remember that task_id is a required named argument for operators.

The correct way to instantiate an operator in a DAG is to provide task_id as a named argument and assign the dag. Option D does this correctly.

❓ Troubleshoot

advanced

Error raised when execute method is missing

What error will Airflow raise if a custom operator class does not implement the execute method?

Apache Airflow

from airflow.models import BaseOperator

class BrokenOperator(BaseOperator):
    pass

# When Airflow tries to run BrokenOperator, what happens?

A AttributeError: 'BrokenOperator' object has no attribute 'execute'

B TypeError: Can't instantiate abstract class BrokenOperator with abstract methods execute

C RuntimeError: execute method missing

D No error, runs with default execute

💡 Hint

BaseOperator uses Python's abstract base class mechanism to enforce execute implementation.

BaseOperator declares execute as an abstract method. Not implementing it causes a TypeError when instantiating the class.

🔀 Workflow

advanced

Order of execution in DAG with custom operators

Given a DAG with two custom operators op1 and op2 where op2 depends on op1, what is the correct order of execution?

Apache Airflow

op1 = MyOperator(task_id='op1', dag=dag)
op2 = MyOperator(task_id='op2', dag=dag)
op1 >> op2

A Both run in parallel

B 2,1

C 1,2

D No guaranteed order

💡 Hint

The >> operator sets task dependencies in Airflow.

op1 >> op2 means op1 must finish before op2 starts, so op1 runs first, then op2.

✅ Best Practice

expert

Best practice for unit testing a custom operator's execute method

Which approach is best for unit testing the execute method of a custom Airflow operator?

A Mock Airflow context and dependencies, then call <code>execute</code> directly and assert results

B Run the operator in a real Airflow environment and check logs manually

C Test only the DAG, not the operator code

D Use print statements inside <code>execute</code> and verify output visually

💡 Hint

Unit tests should be automated and isolated from external systems.

Mocking the Airflow context and dependencies allows isolated, automated testing of the execute method without needing a full Airflow environment.

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