diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_calendar.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_calendar.py index fce45ffe1ef25..250939852ee0a 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_calendar.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_calendar.py @@ -25,6 +25,7 @@ from airflow.utils.session import provide_session from airflow.utils.state import DagRunState +from tests_common.test_utils.asserts import assert_queries_count from tests_common.test_utils.db import clear_db_dags, clear_db_runs pytestmark = pytest.mark.db_test @@ -103,7 +104,8 @@ def teardown_method(self) -> None: ], ) def test_daily_calendar(self, test_client, query_params, result): - response = test_client.get(f"/calendar/{self.DAG_NAME}", params=query_params) + with assert_queries_count(4): + response = test_client.get(f"/calendar/{self.DAG_NAME}", params=query_params) assert response.status_code == 200 body = response.json() print(body) @@ -173,7 +175,8 @@ def test_daily_calendar(self, test_client, query_params, result): ], ) def test_hourly_calendar(self, setup_dag_runs, test_client, query_params, result): - response = test_client.get(f"/calendar/{self.DAG_NAME}", params=query_params) + with assert_queries_count(4): + response = test_client.get(f"/calendar/{self.DAG_NAME}", params=query_params) assert response.status_code == 200 body = response.json()