From ef4e7e1f1dd337b1c215ca0c0e8a7f596e12b467 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Mon, 15 Dec 2025 17:11:41 +0800 Subject: [PATCH] Refactor deprecated SQLA models/test_pool.py --- .pre-commit-config.yaml | 1 + airflow-core/tests/unit/models/test_pool.py | 7 ++++--- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 81e24f2e8a3b3..1a1177b178d82 100644 --- a/.pre-commit-config.yaml +++ b/.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/models/test_pool.py$| ^task_sdk.*\.py$ pass_filenames: true - id: update-supported-versions diff --git a/airflow-core/tests/unit/models/test_pool.py b/airflow-core/tests/unit/models/test_pool.py index fa59d85ffe82d..a275bf2d4b246 100644 --- a/airflow-core/tests/unit/models/test_pool.py +++ b/airflow-core/tests/unit/models/test_pool.py @@ -21,6 +21,7 @@ import pendulum import pytest +from sqlalchemy import func, select from airflow import settings from airflow.exceptions import AirflowException, PoolNotFound @@ -292,7 +293,7 @@ def test_create_pool(self, session): assert pool.slots == 5 assert pool.description == "" assert pool.include_deferred is True - assert session.query(Pool).count() == self.TOTAL_POOL_COUNT + 1 + assert session.scalar(select(func.count()).select_from(Pool)) == self.TOTAL_POOL_COUNT + 1 def test_create_pool_existing(self, session): self.add_pools() @@ -303,13 +304,13 @@ def test_create_pool_existing(self, session): assert pool.slots == 5 assert pool.description == "" assert pool.include_deferred is False - assert session.query(Pool).count() == self.TOTAL_POOL_COUNT + assert session.scalar(select(func.count()).select_from(Pool)) == self.TOTAL_POOL_COUNT def test_delete_pool(self, session): self.add_pools() pool = Pool.delete_pool(name=self.pools[-1].pool) assert pool.pool == self.pools[-1].pool - assert session.query(Pool).count() == self.TOTAL_POOL_COUNT - 1 + assert session.scalar(select(func.count()).select_from(Pool)) == self.TOTAL_POOL_COUNT - 1 def test_delete_pool_non_existing(self): with pytest.raises(PoolNotFound, match="^Pool 'test' doesn't exist$"):