Lee-W commented on code in PR #69075:
URL: https://github.com/apache/airflow/pull/69075#discussion_r3497078204
##########
task-sdk/tests/task_sdk/definitions/test_xcom_arg.py:
##########
@@ -350,3 +350,72 @@ def xcom_get(msg):
states = [run_ti(dag, "pull_one", map_index) for map_index in range(5)]
assert states == [TaskInstanceState.SUCCESS] * 5
assert agg_results == {"a", "b", "c", 1, 2}
+
+
+class TestPlainXComArgResolveMappedGroup:
+ """Resolving a task inside a mapped task group from a task outside that
group.
+
+ Regression tests for #69036 and #48005: the combined return value of a
+ mapped task group must always serialize to a list (one element per
+ expansion), even when the group expanded only once or every expansion
+ returned ``None``. Previously this case was routed through ``xcom_pull``
+ pulling all map indexes, which collapsed a single value to a bare scalar
+ and an empty set of values to ``None``. ``resolve`` stays lazy and serde
+ materialises the sequence only when the value is actually returned.
+ """
+
+ @staticmethod
+ def _make_ti(*, computed):
+ ti = mock.MagicMock()
+ ti._upstream_map_indexes = None
+ ti._cached_template_context = {"expanded_ti_count": 1}
+ ti.run_id = "run-1"
+ ti.get_relevant_upstream_map_indexes.return_value = computed
+ return ti
+
+ @staticmethod
+ def _make_arg():
+ from airflow.sdk.definitions.xcom_arg import PlainXComArg
Review Comment:
I think we can move it to the top level
##########
task-sdk/tests/task_sdk/definitions/test_xcom_arg.py:
##########
@@ -350,3 +350,72 @@ def xcom_get(msg):
states = [run_ti(dag, "pull_one", map_index) for map_index in range(5)]
assert states == [TaskInstanceState.SUCCESS] * 5
assert agg_results == {"a", "b", "c", 1, 2}
+
+
+class TestPlainXComArgResolveMappedGroup:
+ """Resolving a task inside a mapped task group from a task outside that
group.
+
+ Regression tests for #69036 and #48005: the combined return value of a
+ mapped task group must always serialize to a list (one element per
+ expansion), even when the group expanded only once or every expansion
+ returned ``None``. Previously this case was routed through ``xcom_pull``
+ pulling all map indexes, which collapsed a single value to a bare scalar
+ and an empty set of values to ``None``. ``resolve`` stays lazy and serde
+ materialises the sequence only when the value is actually returned.
+ """
+
+ @staticmethod
+ def _make_ti(*, computed):
+ ti = mock.MagicMock()
+ ti._upstream_map_indexes = None
+ ti._cached_template_context = {"expanded_ti_count": 1}
+ ti.run_id = "run-1"
+ ti.get_relevant_upstream_map_indexes.return_value = computed
+ return ti
+
+ @staticmethod
+ def _make_arg():
+ from airflow.sdk.definitions.xcom_arg import PlainXComArg
+
+ operator = mock.MagicMock()
+ operator.is_mapped = False
+ operator.task_id = "do_something"
+ operator.dag_id = "test_dag"
+ operator.get_closest_mapped_task_group.return_value = mock.MagicMock()
+ return PlainXComArg(operator=operator, key="test")
+
+ @pytest.mark.parametrize(
+ ("root", "expected"),
+ [
+ pytest.param(["14"], ["14"], id="single-expansion-stays-a-list"),
+ pytest.param([], [], id="all-none-expansions-give-empty-list"),
+ pytest.param(["a", "b"], ["a", "b"], id="multiple-expansions"),
+ ],
+ )
+ def test_resolve_stays_lazy_and_serializes_as_list(self, root, expected,
mock_supervisor_comms):
+ from airflow.sdk.execution_time.comms import XComSequenceSliceResult
Review Comment:
same here
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
