From 2d1fb47e185f71eb5689165274fd6fd3dc3b7eb9 Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Mon, 12 Aug 2024 00:45:25 +0200 Subject: [PATCH] Simpler task retrieval for taskinstance test The test has been updated for DB isolation but the retrieval of task was not intuitive and it could lead to flaky tests possibly --- tests/models/test_taskinstance.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index c2993e9ce8d95..2d3da6a413e2a 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -1539,8 +1539,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()