fix: sqa deprecations for airflow test

apache · Apr 28, 2024 · a4aaaa9 · a4aaaa9
1 parent 25f901a
commit a4aaaa9
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 4 deletions.
diff --git a/tests/api_connexion/schemas/test_task_instance_schema.py b/tests/api_connexion/schemas/test_task_instance_schema.py
@@ -65,6 +65,7 @@ def set_attrs(self, session, dag_maker):
 
     def test_task_instance_schema_without_sla_and_rendered(self, session):
         ti = TI(task=self.task, **self.default_ti_init)
+        session.add(ti)
         for key, value in self.default_ti_extras.items():
             setattr(ti, key, value)
         serialized_ti = task_instance_schema.dump((ti, None, None))
@@ -109,6 +110,7 @@ def test_task_instance_schema_with_sla_and_rendered(self, session):
         session.add(sla_miss)
         session.flush()
         ti = TI(task=self.task, **self.default_ti_init)
+        session.add(ti)
         for key, value in self.default_ti_extras.items():
             setattr(ti, key, value)
         self.task.template_fields = ["partitions"]

diff --git a/tests/api_experimental/common/test_mark_tasks.py b/tests/api_experimental/common/test_mark_tasks.py
@@ -21,7 +21,7 @@
 from typing import Callable
 
 import pytest
-from sqlalchemy.orm import eagerload
+from sqlalchemy.orm import joinedload
 
 from airflow import models
 from airflow.api.common.mark_tasks import (
@@ -134,7 +134,7 @@ def snapshot_state(dag, execution_dates):
             return (
                 session.query(TI)
                 .join(TI.dag_run)
-                .options(eagerload(TI.dag_run))
+                .options(joinedload(TI.dag_run))
                 .filter(TI.dag_id == dag.dag_id, DR.execution_date.in_(execution_dates))
                 .all()
             )

diff --git a/tests/serialization/test_pydantic_models.py b/tests/serialization/test_pydantic_models.py
@@ -159,8 +159,9 @@ def test_serializing_pydantic_dataset_event(session, create_task_instance, creat
     ds2_event_1 = DatasetEvent(dataset_id=2)
     ds2_event_2 = DatasetEvent(dataset_id=2)
 
-    DagScheduleDatasetReference(dag_id=dag.dag_id, dataset=ds1)
-    TaskOutletDatasetReference(task_id=task1.task_id, dag_id=dag.dag_id, dataset=ds1)
+    dag_ds_ref = DagScheduleDatasetReference(dag_id=dag.dag_id, dataset=ds1)
+    task_ds_ref = TaskOutletDatasetReference(task_id=task1.task_id, dag_id=dag.dag_id, dataset=ds1)
+    session.add_all([dag_ds_ref, task_ds_ref])
 
     dr.consumed_dataset_events.append(ds1_event)
     dr.consumed_dataset_events.append(ds2_event_1)

diff --git a/tests/utils/test_log_handlers.py b/tests/utils/test_log_handlers.py
@@ -493,6 +493,7 @@ def test_set_context_trigger(self, create_dummy_dag, dag_maker, is_a_trigger, se
             job = Job()
             t = Trigger("", {})
             t.triggerer_job = job
+            session.add(t)
             ti.triggerer = t
             t.task_instance = ti
         h = FileTaskHandler(base_log_folder=os.fspath(tmp_path))