Fix xcom_pull for task_ids=None (apache#47407)

Author: insomnes

task-sdk/src/airflow/sdk/execution_time/task_runner.py

@@ -292,7 +292,7 @@ def xcom_pull(
 
         if task_ids is None:
             # default to the current task if not provided
-            task_ids = self.task_id
+            task_ids = [self.task_id]
         elif isinstance(task_ids, str):
             task_ids = [task_ids]
         if isinstance(map_indexes, ArgNotSet):

task-sdk/tests/task_sdk/execution_time/test_task_runner.py

@@ -21,6 +21,7 @@
 import json
 import os
 import uuid
+from collections.abc import Iterable
 from datetime import datetime, timedelta
 from pathlib import Path
 from socket import socketpair
@@ -93,6 +94,7 @@
 )
 from airflow.utils import timezone
 from airflow.utils.state import TaskInstanceState
+from airflow.utils.types import NOTSET, ArgNotSet
 
 from tests_common.test_utils.mock_operators import AirflowLink
 
@@ -1134,28 +1136,37 @@ def test_get_variable_from_context(
             "push_task",
             ["push_task1", "push_task2"],
             {"push_task1", "push_task2"},
+            None,
+            NOTSET,
         ],
     )
     def test_xcom_pull(self, create_runtime_ti, mock_supervisor_comms, spy_agency, task_ids):
         """Test that a task pulls the expected XCom value if it exists."""
 
         class CustomOperator(BaseOperator):
             def execute(self, context):
-                value = context["ti"].xcom_pull(task_ids=task_ids, key="key")
+                if isinstance(task_ids, ArgNotSet):
+                    value = context["ti"].xcom_pull(key="key")
+                else:
+                    value = context["ti"].xcom_pull(task_ids=task_ids, key="key")
                 print(f"Pulled XCom Value: {value}")
 
-        task = CustomOperator(task_id="pull_task")
+        test_task_id = "pull_task"
+        task = CustomOperator(task_id=test_task_id)
 
         runtime_ti = create_runtime_ti(task=task)
 
         mock_supervisor_comms.get_message.return_value = XComResult(key="key", value='"value"')
 
         run(runtime_ti, log=mock.MagicMock())
 
-        if isinstance(task_ids, str):
+        if not isinstance(task_ids, Iterable) or isinstance(task_ids, str):
             task_ids = [task_ids]
 
         for task_id in task_ids:
+            # Without task_ids (or None) expected behavior is to pull with calling task_id
+            if task_id is None or isinstance(task_id, ArgNotSet):
+                task_id = test_task_id
             mock_supervisor_comms.send_request.assert_any_call(
                 log=mock.ANY,
                 msg=GetXCom(