How to test on_failure_callback Airflow operators

Viewed 706

I have a scenario where I want to make an integration test for operators called through on_failure_callback in an Airflow DAG.

A minimal example of this DAG is as follows:

def failure_callback(context):
    # CustomOperator in this case links to an external K8s service
    handle_failure = CustomOperator(
        task_id="handle_failure",
        timestamp=context["ts"]
    )

    handle_failure.execute(context=context)

args = {
    "catchup": False,
    "retries": 3,
    "retry_delay": timedelta(seconds=30),
    "start_date": START_DATE,
    "on_failure_callback": failure_callback,
}

with DAG("foo", schedule_interval=None, default_args=args) as dag:
    
    task_to_fail = SomeOperator()

My first thought for testing would be to run task_to_fail, let it fail, and validate the outcome of the failure_callback with some other process, attempt below:

import pytest
from airflow.models import DagBag, TaskInstance
from dateutil import parser

@pytest.fixture
def foo_dag():
    dag_id = "foo"
    dag_bag = DagBag("dags")
    return dag_bag.dags[dag_id]

@pytest.mark.integration
def test_task_to_fail(foo_dag):
    execution_date = parser.parse("2000-01-01T00:00+00:00")
    task_id = "task_to_fail"

    task = foo_dag.get_task(task_id=task_id)
    task_instance = TaskInstance(task, execution_date)

    with pytest.raises(Exception):
        task_instance.run(ignore_task_deps=True, ignore_ti_state=True, test_mode=True)

    assert "INSERT DESIRED OUTCOME OF `failure_callback` HERE"

The issue I'm having is that it doesn't appear that failure_callback is being called when running pytest. I suspect this is due to how TaskInstance is being called (i.e not running the on_failure_callback, but am not sure.

My questions:

  1. Is this the correct way to validate the behavior of this callback? If not, how should this be handled?
  2. Upstream of the task_to_fail task, there are many expensive operations that I want to avoid running during tests. Is it possible to have a full-run of a DAG, executed with pytest, starting from a particular task (in this case, task_to_fail?
0 Answers
Related