diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 939039b0827fb..321cd127d426a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -425,6 +425,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$| ^airflow-core/tests/unit/utils/test_db_cleanup.py$| ^dev/airflow_perf/scheduler_dag_execution_timing.py$| ^providers/openlineage/.*\.py$| 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$"):