Refactor deprecated SQLA test_db_cleanup.py

.pre-commit-config.yaml

@@ -413,6 +413,7 @@ repos:
             ^airflow-ctl.*\.py$|
             ^airflow-core/src/airflow/models/.*\.py$|
             ^airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py$|
+            ^airflow-core/tests/unit/utils/test_db_cleanup.py$|
             ^dev/airflow_perf/scheduler_dag_execution_timing.py$|
             ^providers/openlineage/.*\.py$|
             ^task_sdk.*\.py$

airflow-core/tests/unit/utils/test_db_cleanup.py

@@ -26,7 +26,7 @@
 
 import pendulum
 import pytest
-from sqlalchemy import inspect, text
+from sqlalchemy import func, inspect, select, text
 from sqlalchemy.exc import OperationalError, SQLAlchemyError
 from sqlalchemy.ext.declarative import DeclarativeMeta
 
@@ -301,11 +301,11 @@ def test__cleanup_table(self, table_name, date_add_kwargs, expected_to_delete, r
             )
             model = config_dict[table_name].orm_model
             expected_remaining = num_tis - expected_to_delete
-            assert len(session.query(model).all()) == expected_remaining
+            assert session.scalar(select(func.count()).select_from(model)) == expected_remaining
             if model.name == "task_instance":
-                assert len(session.query(DagRun).all()) == num_tis
+                assert session.scalar(select(func.count()).select_from(DagRun)) == num_tis
             elif model.name == "dag_run":
-                assert len(session.query(TaskInstance).all()) == expected_remaining
+                assert session.scalar(select(func.count()).select_from(TaskInstance)) == expected_remaining
             else:
                 raise Exception("unexpected")
 
@@ -384,7 +384,7 @@ def test__skip_archive(self, skip_archive, expected_archives):
                 skip_archive=skip_archive,
             )
             model = config_dict["dag_run"].orm_model
-            assert len(session.query(model).all()) == 5
+            assert session.scalar(select(func.count()).select_from(model)) == 5
             assert len(_get_archived_table_names(["dag_run"], session)) == expected_archives
 
     @patch("airflow.utils.db.reflect_tables")