Fix test mock using relevant interface

Author: uranusjr

airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py

@@ -142,20 +142,20 @@ def get_xcom(
     params: Annotated[GetXcomFilterParams, Query()],
 ) -> XComResponse:
     """Get an Airflow XCom from database - not other XCom Backends."""
-    # The xcom_query allows no map_index to be passed. This endpoint should always return just a single item,
-    # so we override that query value
     xcom_query = XComModel.get_many(
         run_id=run_id,
         key=key,
         task_ids=task_id,
         dag_ids=dag_id,
-        map_indexes=None if params.offset is not None else params.map_index,
         include_prior_dates=params.include_prior_dates,
-        latest_first=(params.offset is None or params.offset >= 0),
         session=session,
     )
     if params.offset is not None:
-        xcom_query = xcom_query.offset(abs(params.offset))
+        xcom_query = xcom_query.filter(XComModel.value.is_not(None)).order_by(None)
+        if params.offset >= 0:
+            xcom_query = xcom_query.order_by(XComModel.map_index.asc()).offset(params.offset)
+        else:
+            xcom_query = xcom_query.order_by(XComModel.map_index.desc()).offset(-1 - params.offset)
     else:
         xcom_query = xcom_query.filter(XComModel.map_index == params.map_index)
 
@@ -166,13 +166,19 @@ def get_xcom(
     # performance hits from retrieving large data files into the API server.
     result = xcom_query.limit(1).first()
     if result is None:
-        map_index = params.map_index
+        if params.offset is None:
+            message = (
+                f"XCom with {key=} map_index={params.map_index} not found for "
+                f"task {task_id!r} in DAG run {run_id!r} of {dag_id!r}"
+            )
+        else:
+            message = (
+                f"XCom with {key=} offset={params.offset} not found for "
+                f"task {task_id!r} in DAG run {run_id!r} of {dag_id!r}"
+            )
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
-            detail={
-                "reason": "not_found",
-                "message": f"XCom with {key=} {map_index=} not found for task {task_id!r} in DAG run {run_id!r} of {dag_id!r}",
-            },
+            detail={"reason": "not_found", "message": message},
         )
 
     return XComResponse(key=key, value=result.value)

airflow-core/src/airflow/models/xcom.py

@@ -252,7 +252,6 @@ def get_many(
         map_indexes: int | Iterable[int] | None = None,
         include_prior_dates: bool = False,
         limit: int | None = None,
-        latest_first: bool = True,
         session: Session = NEW_SESSION,
     ) -> Query:
         """
@@ -275,9 +274,7 @@ def get_many(
             returned regardless of the run it belongs to.
         :param session: Database session. If not given, a new session will be
             created for this function.
-        :param limit: Limiting returning XComs.
-        :param latest_first: If *True* (default), returning XComs are ordered
-            latest-first. Otherwise earlier XComs are returned first.
+        :param limit: Limiting returning XComs
         """
         from airflow.models.dagrun import DagRun
 
@@ -321,10 +318,7 @@ def get_many(
         else:
             query = query.filter(cls.run_id == run_id)
 
-        if latest_first:
-            query = query.order_by(DagRun.logical_date.desc(), cls.timestamp.desc())
-        else:
-            query = query.order_by(DagRun.logical_date.asc(), cls.timestamp.asc())
+        query = query.order_by(DagRun.logical_date.desc(), cls.timestamp.desc())
         if limit:
             return query.limit(limit)
         return query

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py

@@ -29,6 +29,7 @@
 from airflow.models.dagrun import DagRun
 from airflow.models.taskmap import TaskMap
 from airflow.models.xcom import XComModel
+from airflow.providers.standard.operators.empty import EmptyOperator
 from airflow.serialization.serde import deserialize, serialize
 from airflow.utils.session import create_session
 
@@ -130,6 +131,86 @@ def test_xcom_access_denied(self, client, caplog):
         }
         assert any(msg.startswith("Checking read XCom access") for msg in caplog.messages)
 
+    @pytest.mark.parametrize(
+        "offset, expected_status, expected_json",
+        [
+            pytest.param(
+                -4,
+                404,
+                {
+                    "detail": {
+                        "reason": "not_found",
+                        "message": (
+                            "XCom with key='xcom_1' offset=-4 not found "
+                            "for task 'task' in DAG run 'runid' of 'dag'"
+                        ),
+                    },
+                },
+                id="-4",
+            ),
+            pytest.param(-3, 200, {"key": "xcom_1", "value": "f"}, id="-3"),
+            pytest.param(-2, 200, {"key": "xcom_1", "value": "o"}, id="-2"),
+            pytest.param(-1, 200, {"key": "xcom_1", "value": "b"}, id="-1"),
+            pytest.param(0, 200, {"key": "xcom_1", "value": "f"}, id="0"),
+            pytest.param(1, 200, {"key": "xcom_1", "value": "o"}, id="1"),
+            pytest.param(2, 200, {"key": "xcom_1", "value": "b"}, id="2"),
+            pytest.param(
+                3,
+                404,
+                {
+                    "detail": {
+                        "reason": "not_found",
+                        "message": (
+                            "XCom with key='xcom_1' offset=3 not found "
+                            "for task 'task' in DAG run 'runid' of 'dag'"
+                        ),
+                    },
+                },
+                id="3",
+            ),
+        ],
+    )
+    def test_xcom_get_with_offset(
+        self,
+        client,
+        dag_maker,
+        session,
+        offset,
+        expected_status,
+        expected_json,
+    ):
+        xcom_values = ["f", None, "o", "b"]
+
+        class MyOperator(EmptyOperator):
+            def __init__(self, *, x, **kwargs):
+                super().__init__(**kwargs)
+                self.x = x
+
+        with dag_maker(dag_id="dag"):
+            MyOperator.partial(task_id="task").expand(x=xcom_values)
+        dag_run = dag_maker.create_dagrun(run_id="runid")
+        tis = {ti.map_index: ti for ti in dag_run.task_instances}
+
+        for map_index, db_value in enumerate(xcom_values):
+            if db_value is None:  # We don't put None to XCom.
+                continue
+            ti = tis[map_index]
+            x = XComModel(
+                key="xcom_1",
+                value=db_value,
+                dag_run_id=ti.dag_run.id,
+                run_id=ti.run_id,
+                task_id=ti.task_id,
+                dag_id=ti.dag_id,
+                map_index=map_index,
+            )
+            session.add(x)
+        session.commit()
+
+        response = client.get(f"/execution/xcoms/dag/runid/task/xcom_1?offset={offset}")
+        assert response.status_code == expected_status
+        assert response.json() == expected_json
+
 
 class TestXComsSetEndpoint:
     @pytest.mark.parametrize(

airflow-core/tests/unit/models/test_mappedoperator.py

@@ -32,7 +32,7 @@
 from airflow.models.taskmap import TaskMap
 from airflow.providers.standard.operators.python import PythonOperator
 from airflow.sdk import setup, task, task_group, teardown
-from airflow.sdk.execution_time.comms import XComCountResponse
+from airflow.sdk.execution_time.comms import XComCountResponse, XComResult
 from airflow.utils.state import TaskInstanceState
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
@@ -1270,8 +1270,16 @@ def my_teardown(val):
         ) as supervisor_comms:
             # TODO: TaskSDK: this is a bit of a hack that we need to stub this at all. `dag.test()` should
             # really work without this!
-            supervisor_comms.get_message.return_value = XComCountResponse(len=3)
+            supervisor_comms.get_message.side_effect = [
+                XComCountResponse(len=3),
+                XComResult(key="return_value", value=1),
+                XComCountResponse(len=3),
+                XComResult(key="return_value", value=2),
+                XComCountResponse(len=3),
+                XComResult(key="return_value", value=3),
+            ]
             dr = dag.test()
+            assert supervisor_comms.get_message.call_count == 6
         states = self.get_states(dr)
         expected = {
             "tg_1.my_pre_setup": "success",

task-sdk/src/airflow/sdk/definitions/xcom_arg.py

@@ -337,7 +337,7 @@ def resolve(self, context: Mapping[str, Any]) -> Any:
         task_id = self.operator.task_id
 
         if self.operator.is_mapped:
-            return LazyXComSequence[Any](xcom_arg=self, ti=ti)
+            return LazyXComSequence(xcom_arg=self, ti=ti)
         tg = self.operator.get_closest_mapped_task_group()
         result = None
         if tg is None:

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

@@ -0,0 +1,131 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from unittest.mock import ANY, Mock, call
+
+import pytest
+
+from airflow.sdk.exceptions import ErrorType
+from airflow.sdk.execution_time.comms import (
+    ErrorResponse,
+    GetXComCount,
+    GetXComSequenceItem,
+    XComCountResponse,
+    XComResult,
+)
+from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
+
+
+@pytest.fixture
+def mock_operator():
+    return Mock(spec=["dag_id", "task_id"], dag_id="dag", task_id="task")
+
+
+@pytest.fixture
+def mock_xcom_arg(mock_operator):
+    return Mock(spec=["operator", "key"], operator=mock_operator, key="return_value")
+
+
+@pytest.fixture
+def mock_ti():
+    return Mock(spec=["run_id"], run_id="run")
+
+
+@pytest.fixture
+def lazy_sequence(mock_xcom_arg, mock_ti):
+    return LazyXComSequence(mock_xcom_arg, mock_ti)
+
+
+def test_len(mock_supervisor_comms, lazy_sequence):
+    mock_supervisor_comms.get_message.return_value = XComCountResponse(len=3)
+    assert len(lazy_sequence) == 3
+    assert mock_supervisor_comms.send_request.mock_calls == [
+        call(log=ANY, msg=GetXComCount(key="return_value", dag_id="dag", task_id="task", run_id="run")),
+    ]
+
+
+def test_iter(mock_supervisor_comms, lazy_sequence):
+    it = iter(lazy_sequence)
+
+    mock_supervisor_comms.get_message.side_effect = [
+        XComResult(key="return_value", value="f"),
+        ErrorResponse(error=ErrorType.XCOM_NOT_FOUND, detail={"oops": "sorry!"}),
+    ]
+    assert list(it) == ["f"]
+    assert mock_supervisor_comms.send_request.mock_calls == [
+        call(
+            log=ANY,
+            msg=GetXComSequenceItem(
+                key="return_value",
+                dag_id="dag",
+                task_id="task",
+                run_id="run",
+                offset=0,
+            ),
+        ),
+        call(
+            log=ANY,
+            msg=GetXComSequenceItem(
+                key="return_value",
+                dag_id="dag",
+                task_id="task",
+                run_id="run",
+                offset=1,
+            ),
+        ),
+    ]
+
+
+def test_getitem_index(mock_supervisor_comms, lazy_sequence):
+    mock_supervisor_comms.get_message.return_value = XComResult(key="return_value", value="f")
+    assert lazy_sequence[4] == "f"
+    assert mock_supervisor_comms.send_request.mock_calls == [
+        call(
+            log=ANY,
+            msg=GetXComSequenceItem(
+                key="return_value",
+                dag_id="dag",
+                task_id="task",
+                run_id="run",
+                offset=4,
+            ),
+        ),
+    ]
+
+
+def test_getitem_indexerror(mock_supervisor_comms, lazy_sequence):
+    mock_supervisor_comms.get_message.return_value = ErrorResponse(
+        error=ErrorType.XCOM_NOT_FOUND,
+        detail={"oops": "sorry!"},
+    )
+    with pytest.raises(IndexError) as ctx:
+        lazy_sequence[4]
+    assert ctx.value.args == (4,)
+    assert mock_supervisor_comms.send_request.mock_calls == [
+        call(
+            log=ANY,
+            msg=GetXComSequenceItem(
+                key="return_value",
+                dag_id="dag",
+                task_id="task",
+                run_id="run",
+                offset=4,
+            ),
+        ),
+    ]