fix tests

Author: dstandish

airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py

@@ -719,7 +719,6 @@ def post_clear_task_instances(
         clear_task_instances(
             task_instances,
             session,
-            dag,
             DagRunState.QUEUED if reset_dag_runs else False,
         )
 

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

@@ -180,12 +180,6 @@ def teardown_method(self) -> None:
         clear_db_dags()
         clear_db_assets()
 
-    @staticmethod
-    def _clean_up(dag_id: str):
-        with create_session() as session:
-            session.query(DagRun).filter(DagRun.dag_id == dag_id).delete(synchronize_session=False)
-            session.query(TI).filter(TI.dag_id == dag_id).delete(synchronize_session=False)
-
     @staticmethod
     def _occur_before(a, b, list_):
         """
@@ -977,8 +971,6 @@ def add_failed_dag_run(dag, id, logical_date):
         )
         add_failed_dag_run(dag, "2", TEST_DATE + timedelta(days=1))
         assert dag.get_is_paused()
-        dag.clear()
-        self._clean_up(dag_id)
 
     def test_dag_is_deactivated_upon_dagfile_deletion(self, dag_maker):
         dag_id = "old_existing_dag"
@@ -1038,8 +1030,6 @@ def test_schedule_dag_no_previous_runs(self):
         )
         assert dag_run.state == State.RUNNING
         assert dag_run.run_type != DagRunType.MANUAL
-        dag.clear()
-        self._clean_up(dag_id)
 
     @patch("airflow.models.dag.Stats")
     def test_dag_handle_callback_crash(self, mock_stats):
@@ -1080,9 +1070,6 @@ def test_dag_handle_callback_crash(self, mock_stats):
             tags={"dag_id": "test_dag_callback_crash"},
         )
 
-        dag.clear()
-        self._clean_up(dag_id)
-
     def test_dag_handle_callback_with_removed_task(self, dag_maker, session):
         """
         Tests avoid crashes when a removed task is the last one in the list of task instance
@@ -1118,9 +1105,6 @@ def test_dag_handle_callback_with_removed_task(self, dag_maker, session):
             dag.handle_callback(dag_run, success=True)
             dag.handle_callback(dag_run, success=False)
 
-        dag.clear()
-        self._clean_up(dag_id)
-
     @pytest.mark.parametrize("catchup,expected_next_dagrun", [(True, DEFAULT_DATE), (False, None)])
     def test_next_dagrun_after_fake_scheduled_previous(self, catchup, expected_next_dagrun):
         """
@@ -1158,8 +1142,6 @@ def test_next_dagrun_after_fake_scheduled_previous(self, catchup, expected_next_
             assert model.next_dagrun == expected_next_dagrun
             assert model.next_dagrun_create_after == expected_next_dagrun + delta
 
-        self._clean_up(dag_id)
-
     def test_schedule_dag_once(self):
         """
         Tests scheduling a dag scheduled for @once - should be scheduled the first time
@@ -1188,7 +1170,6 @@ def test_schedule_dag_once(self):
 
         assert model.next_dagrun is None
         assert model.next_dagrun_create_after is None
-        self._clean_up(dag_id)
 
     def test_fractional_seconds(self):
         """
@@ -1213,7 +1194,6 @@ def test_fractional_seconds(self):
 
         assert start_date == run.logical_date, "dag run logical_date loses precision"
         assert start_date == run.start_date, "dag run start_date loses precision "
-        self._clean_up(dag_id)
 
     def test_rich_comparison_ops(self):
         test_dag_id = "test_rich_comparison_ops"
@@ -1397,46 +1377,39 @@ def test_dag_add_task_sets_default_task_group(self):
         assert dag.get_task("task_group.task_with_task_group") == task_with_task_group
 
     @pytest.mark.parametrize("dag_run_state", [DagRunState.QUEUED, DagRunState.RUNNING])
-    def test_clear_set_dagrun_state(self, dag_run_state):
+    @pytest.mark.need_serialized_dag
+    def test_clear_set_dagrun_state(self, dag_run_state, dag_maker, session):
         dag_id = "test_clear_set_dagrun_state"
-        self._clean_up(dag_id)
-        task_id = "t1"
-        dag = DAG(dag_id, schedule=None, start_date=DEFAULT_DATE, max_active_runs=1)
-        t_1 = EmptyOperator(task_id=task_id, dag=dag)
 
-        session = settings.Session()
-        dagrun_1 = _create_dagrun(
-            dag,
+        with dag_maker(dag_id, start_date=DEFAULT_DATE, max_active_runs=1) as dag:
+            task_id = "t1"
+            EmptyOperator(task_id=task_id)
+
+        dr = dag_maker.create_dagrun(
             run_type=DagRunType.BACKFILL_JOB,
             state=State.FAILED,
             start_date=DEFAULT_DATE,
             logical_date=DEFAULT_DATE,
-            data_interval=(DEFAULT_DATE, DEFAULT_DATE),
+            session=session,
         )
-        session.merge(dagrun_1)
-
-        task_instance_1 = TI(t_1, run_id=dagrun_1.run_id, state=State.RUNNING)
-        task_instance_1.refresh_from_db()
-        session.merge(task_instance_1)
         session.commit()
+        session.refresh(dr)
+        assert dr.state == "failed"
 
         dag.clear(
             start_date=DEFAULT_DATE,
             end_date=DEFAULT_DATE + datetime.timedelta(days=1),
             dag_run_state=dag_run_state,
             session=session,
         )
-        dagruns = session.query(DagRun).filter(DagRun.dag_id == dag_id).all()
-
-        assert len(dagruns) == 1
-        dagrun: DagRun = dagruns[0]
-        assert dagrun.state == dag_run_state
+        session.refresh(dr)
+        assert dr.state == dag_run_state
 
     @pytest.mark.parametrize("dag_run_state", [DagRunState.QUEUED, DagRunState.RUNNING])
     @pytest.mark.need_serialized_dag
     def test_clear_set_dagrun_state_for_mapped_task(self, dag_maker, dag_run_state):
         dag_id = "test_clear_set_dagrun_state"
-        self._clean_up(dag_id)
+
         task_id = "t1"
 
         with dag_maker(dag_id, schedule=None, start_date=DEFAULT_DATE, max_active_runs=1) as dag:
@@ -1611,32 +1584,37 @@ def test_clear_dag(
         self,
         ti_state_begin: TaskInstanceState | None,
         ti_state_end: TaskInstanceState | None,
+        dag_maker,
+        session,
     ):
         dag_id = "test_clear_dag"
-        self._clean_up(dag_id)
+
         task_id = "t1"
-        dag = DAG(dag_id, schedule=None, start_date=DEFAULT_DATE, max_active_runs=1)
-        _ = EmptyOperator(task_id=task_id, dag=dag)
+        with dag_maker(
+            dag_id,
+            schedule=None,
+            start_date=DEFAULT_DATE,
+            max_active_runs=1,
+            serialized=True,
+        ) as dag:
+            EmptyOperator(task_id=task_id)
 
         session = settings.Session()  # type: ignore
-        dagrun_1 = dag.create_dagrun(
+        dagrun_1 = dag_maker.create_dagrun(
             run_id="backfill",
             run_type=DagRunType.BACKFILL_JOB,
             state=DagRunState.RUNNING,
             start_date=DEFAULT_DATE,
             logical_date=DEFAULT_DATE,
-            data_interval=(DEFAULT_DATE, DEFAULT_DATE),
-            run_after=DEFAULT_DATE,
-            triggered_by=DagRunTriggeredByType.TEST,
+            # triggered_by=DagRunTriggeredByType.TEST,
+            session=session,
         )
-        session.merge(dagrun_1)
 
-        task_instance_1 = dagrun_1.get_task_instance(task_id)
+        task_instance_1 = dagrun_1.get_task_instance(task_id, session=session)
         if TYPE_CHECKING:
             assert task_instance_1
         task_instance_1.state = ti_state_begin
         task_instance_1.job_id = 123
-        session.merge(task_instance_1)
         session.commit()
 
         dag.clear(
@@ -1650,7 +1628,6 @@ def test_clear_dag(
         assert len(task_instances) == 1
         task_instance: TI = task_instances[0]
         assert task_instance.state == ti_state_end
-        self._clean_up(dag_id)
 
     def test_next_dagrun_info_once(self):
         dag = DAG("test_scheduler_dagrun_once", start_date=timezone.datetime(2015, 1, 1), schedule="@once")
@@ -2508,7 +2485,12 @@ def test_count_number_queries(self, tasks_count):
 def test_set_task_instance_state(run_id, session, dag_maker):
     """Test that set_task_instance_state updates the TaskInstance state and clear downstream failed"""
     start_date = datetime_tz(2020, 1, 1)
-    with dag_maker("test_set_task_instance_state", start_date=start_date, session=session) as dag:
+    with dag_maker(
+        "test_set_task_instance_state",
+        start_date=start_date,
+        session=session,
+        serialized=True,
+    ) as dag:
         task_1 = EmptyOperator(task_id="task_1")
         task_2 = EmptyOperator(task_id="task_2")
         task_3 = EmptyOperator(task_id="task_3")
@@ -2646,7 +2628,12 @@ def consumer(value):
 def test_set_task_group_state(session, dag_maker):
     """Test that set_task_group_state updates the TaskGroup state and clear downstream failed"""
     start_date = datetime_tz(2020, 1, 1)
-    with dag_maker("test_set_task_group_state", start_date=start_date, session=session) as dag:
+    with dag_maker(
+        "test_set_task_group_state",
+        start_date=start_date,
+        session=session,
+        serialized=True,
+    ) as dag:
         start = EmptyOperator(task_id="start")
 
         with TaskGroup("section_1", tooltip="Tasks for section_1") as section_1:

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

@@ -397,19 +397,24 @@ def test_expand_mapped_task_instance_with_named_index(
     expected_rendered_names,
 ) -> None:
     """Test that the correct number of downstream tasks are generated when mapping with an XComArg"""
-    with dag_maker("test-dag", session=session, start_date=DEFAULT_DATE):
+    dag_id = "test_dag_12345"
+    with dag_maker(
+        dag_id=dag_id,
+        start_date=DEFAULT_DATE,
+        serialized=True,
+    ):
         create_mapped_task(task_id="task1", map_names=["a", "b"], template=template)
 
-    dr = dag_maker.create_dagrun()
-    tis = dr.get_task_instances()
+    dr = dag_maker.create_dagrun(session=session)
+    tis = dr.get_task_instances(session=session)
     for ti in tis:
         ti.run()
     session.flush()
 
     indices = session.scalars(
         select(TaskInstance.rendered_map_index)
         .where(
-            TaskInstance.dag_id == "test-dag",
+            TaskInstance.dag_id == dag_id,
             TaskInstance.task_id == "task1",
             TaskInstance.run_id == dr.run_id,
         )