Testing custom operators in Apache Airflow - Cheat Sheet & Quick Revision

Recall & Review

beginner

What is the main purpose of testing custom operators in Airflow?

To ensure that the custom operator behaves as expected and integrates correctly within Airflow workflows before deploying to production.

beginner

Which Python module is commonly used to write unit tests for Airflow custom operators?

The unittest module is commonly used for writing unit tests in Python, including tests for Airflow custom operators.

intermediate

Why should you mock Airflow's context when testing custom operators?

Mocking Airflow's context allows you to simulate the runtime environment and test the operator's logic without running a full Airflow scheduler or executor.

intermediate

What is a simple way to test the execute() method of a custom operator?

Create a test case that calls the execute() method with a mocked context and asserts expected outcomes or side effects.

advanced

How can you verify that your custom operator handles errors correctly during testing?

By writing tests that simulate error conditions and asserting that the operator raises the expected exceptions or handles errors gracefully.

What is the first step when testing a custom Airflow operator?

A Mock the Airflow context

B Run the full Airflow scheduler

C Deploy the operator to production

D Write a DAG using the operator

Mocking the Airflow context allows isolated testing of the operator without needing the full Airflow environment.

Which method in a custom operator is usually tested to check its main functionality?

A execute()

B run()

C start()

D __init__()

The execute() method contains the core logic of the operator and is the main focus of testing.

Why is it important to mock external services in custom operator tests?

A To make tests fail intentionally

B To speed up tests and avoid dependency on external systems

C To deploy faster

D To avoid writing tests

Mocking external services ensures tests run quickly and reliably without depending on outside systems.

Which Python library is commonly used for mocking in tests?

A requests

B pandas

C flask

D mock

The mock library (built-in as unittest.mock) is used to replace parts of the system under test with mock objects.

What should a test assert after calling the execute() method of a custom operator?

A That the operator returns None

B That the operator deletes DAG files

C That the operator performs expected actions or returns expected results

D That the operator runs forever

Tests should verify that the operator performs its intended actions or returns expected results.

Describe the steps to write a unit test for a custom Airflow operator's execute() method.

Think about isolating the operator logic and simulating Airflow's runtime environment.

Explain why mocking is important when testing custom operators in Airflow.

Consider how tests can run without external dependencies.

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