dstandish commented on code in PR #69565:
URL: https://github.com/apache/airflow/pull/69565#discussion_r3545766427
##########
airflow-core/tests/unit/models/test_taskinstance.py:
##########
@@ -3291,6 +3291,159 @@ def show(value):
dag_maker.run_ti(ti.task_id, dag_run=dag_run,
map_index=ti.map_index, session=session)
assert outputs == expected_outputs
+ def test_map_xcom_wide_batched_expand(self, dag_maker, session):
+ """Wide XCom-driven expand goes through the batched add_all()/flush()
path.
+
+ Exercises ``TaskMap.expand_mapped_task`` over a 20-element upstream
XCom and
+ asserts the batched expansion creates exactly N mapped TIs with
contiguous
+ ``map_index`` 0..N-1, the expected ``None`` (schedulable) state, and
that the
+ returned instances are usable: they keep their ``.task`` (no merge()
that drops
+ it), are attached to the session, and have ``dag_run`` primed so a
later
+ ``ti.get_dagrun()`` -- as dependency evaluation makes per index -- is
a cache hit
+ rather than an N+1 SELECT.
+ """
+ from sqlalchemy import event
+ from sqlalchemy.orm.base import NO_VALUE
+
+ width = 20
+ upstream_return = list(range(width))
+
+ with dag_maker(dag_id="xcom_wide", session=session, serialized=True)
as dag:
+
+ @dag.task
+ def emit():
+ return upstream_return
+
+ @dag.task
+ def show(value):
+ return value
+
+ show.expand(value=emit())
+
+ dag_run = dag_maker.create_dagrun()
+ emit_ti = dag_run.get_task_instance("emit", session=session)
+ emit_ti.refresh_from_task(dag_maker.serialized_dag.get_task("emit"))
+ dag_maker.run_ti(emit_ti.task_id, dag_run=dag_run, session=session)
+
+ show_task = dag_maker.serialized_dag.get_task("show")
+ mapped_tis, max_map_index = TaskMap.expand_mapped_task(show_task,
dag_run.run_id, session=session)
Review Comment:
added a query count test for this in commit
dc2fbf544594d93b5765ec1149053f05ba65a25d
--
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]
