Simpler task retrieval for taskinstance test (apache#41389) (apache#4…

…1953) The test has been updated for DB isolation but the retrieval of task was not intuitive and it could lead to flaky tests possibly (cherry picked from commit f25adf1) Co-authored-by: Jarek Potiuk <[email protected]>
aws-mwaa · Sep 2, 2024 · d7d8814 · d7d8814
1 parent f036f36
commit d7d8814
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
@@ -1540,8 +1540,9 @@ def do_something_else(i):
         monkeypatch.setattr(_UpstreamTIStates, "calculate", lambda *_: upstream_states)
         ti = dr.get_task_instance("do_something_else", session=session)
         ti.map_index = 0
+        base_task = ti.task
         for map_index in range(1, 5):
-            ti = TaskInstance(dr.task_instances[-1].task, run_id=dr.run_id, map_index=map_index)
+            ti = TaskInstance(base_task, run_id=dr.run_id, map_index=map_index)
             session.add(ti)
             ti.dag_run = dr
         session.flush()