What is the correct way to unit test the PythonBranchOperator

[skiedude]

We are currently on airflow 2.4.3 (where this test has worked). In testing what its going to take to update to higher versions, starting with 2.5.0, we hit RecursionError: maximum recursion depth exceeded in comparison in some of our UnitTests that have worked up until now.

Using pytest and following test

    (
        'should_check',
        [{'host1': {'is_physical_device': '0'}, 'host2': {'is_physical_device': '1'}}],
        {'get_info', 'process_servers'}
    ),

]
...
def test_should_check(xcom_patch, task_id, xcom_return_val, expected_downstream, dagbag, create_dag_task_context):

    xcom_patch.return_value = xcom_return_val
    task, context = create_dag_task_context('dag_name', task_id, dagbag)
    task.execute(context)
    xcom_return = task.xcom_pull(task_ids=task_id, key='skipmixin_key', context=context)
    assert set(xcom_return['followed']) == expected_downstream

We use the following helper function to create the task and context to run with task.execute(context)

    def create_dag_task_context():
    """
    Creates a manual dagrun of the given dag, and returns the task and context
    """
    def _create(dag_id, task_id, dagbag):
        dag = dagbag.get_dag(dag_id)
        execution_date = datetime.now()
        dagrun = dag.create_dagrun(
            execution_date=execution_date,
            state=DagRunState.RUNNING,
            run_type=DagRunType.MANUAL
        )
        task = dag.get_task(task_id)
        taskinst = dagrun.get_task_instance(task_id)
        taskinst.task = task
        context = taskinst.get_template_context()
        return task, context
    return _create

The task should_check is a branch task. It can either run the next step get__info or skip it and go straight to process_servers

The official error is

  2025-01-10 12:39:53,638 INFO     /home/airflow/.local/lib/python3.9/site-packages/airflow/operators/python.py:210: in execute
2025-01-10 12:39:53,638 INFO         branch = super().execute(context)
2025-01-10 12:39:53,638 INFO     /home/airflow/.local/lib/python3.9/site-packages/airflow/operators/python.py:173: in execute
2025-01-10 12:39:53,638 INFO         self.op_kwargs = self.determine_kwargs(context)
2025-01-10 12:39:53,639 INFO     /home/airflow/.local/lib/python3.9/site-packages/airflow/models/baseoperator.py:1008: in __setattr__
2025-01-10 12:39:53,639 INFO         self.set_xcomargs_dependencies()
2025-01-10 12:39:53,639 INFO     /home/airflow/.local/lib/python3.9/site-packages/airflow/models/baseoperator.py:1115: in set_xcomargs_dependencies
2025-01-10 12:39:53,639 INFO         XComArg.apply_upstream_relationship(self, arg)
2025-01-10 12:39:53,639 INFO     /home/airflow/.local/lib/python3.9/site-packages/airflow/models/xcom_arg.py:121: in apply_upstream_relationship
2025-01-10 12:39:53,639 INFO         for operator, _ in XComArg.iter_xcom_references(arg):
2025-01-10 12:39:53,639 INFO     /home/airflow/.local/lib/python3.9/site-packages/airflow/models/xcom_arg.py:108: in iter_xcom_references
2025-01-10 12:39:53,639 INFO         yield from XComArg.iter_xcom_references(elem)
2025-01-10 12:39:53,639 INFO     /home/airflow/.local/lib/python3.9/site-packages/airflow/models/xcom_arg.py:111: in iter_xcom_references
2025-01-10 12:39:53,639 INFO         yield from XComArg.iter_xcom_references(getattr(arg, attr))
2025-01-10 12:39:53,639 INFO     /home/airflow/.local/lib/python3.9/site-packages/airflow/models/xcom_arg.py:108: in iter_xcom_references
2025-01-10 12:39:53,640 INFO         yield from XComArg.iter_xcom_references(elem)
2025-01-10 12:39:53,640 INFO     E   RecursionError: maximum recursion depth exceeded in comparison
2025-01-10 12:39:53,640 INFO     !!! Recursion detected (same locals & position)

In looking in the release notes for 2.5.0, this seems to be the cause #27771

In debugging through the code, its pushing the whole context into the op_kwargs key here https://github.com/apache/airflow/blob/2.5.0/airflow/utils/operator_helpers.py#L155-L165

Then when it goes through XcomArg recursion it just endlessly loops through itself.

Any ideas on how to better setup this test?

[skiedude]

The dag itself runs fine, its just the unittest that struggles

[skiedude]

Guess this was a case of the rubber duck. I had read over https://airflow.apache.org/docs/apache-airflow/stable/best-practices.html#unit-tests previously, and tried the suggestion of testing a custom operator. However I realized I tried task.run() and not ti.run()

The following changes got the test working again all the way up to 2.10.4

@@ -27,10 +28,10 @@ def create_dag_task_context():
             state=DagRunState.RUNNING,
             run_type=DagRunType.MANUAL
         )
-        task = dag.get_task(task_id)
-        taskinst = dagrun.get_task_instance(task_id)
-        taskinst.task = task
-        context = taskinst.get_template_context()
-        return task, context
-    return _create
-
+        ti = dagrun.get_task_instance(task_id)
+        ti.task = dag.get_task(task_id)
+        ti.task.trigger_rule = TriggerRule.ALWAYS
+        ti.task.on_success_callback = None
+        context = ti.get_template_context()
+        return ti, context
+    return _create

And in the test

-    task, context = create_dag_task_context('dag_name', task_id, dagbag)
-    task.execute(context)
-    xcom_return = task.xcom_pull(task_ids=task_id, key='skipmixin_key', context=context)
+    ti, context = create_dag_task_context('dag_name', task_id, dagbag)
+    ti.run(ignore_ti_state=True)
+    xcom_return = ti.task.xcom_pull(task_ids=task_id, key='skipmixin_key', context=context)