From b9a5175f1a9b9024e46c5709fc566ea201598c3c Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Tue, 9 Jun 2026 02:57:55 +0100 Subject: [PATCH 1/4] Reuse a session-scoped FastAPI app across api_fastapi tests The test_client / unauthenticated_test_client / unauthorized_test_client fixtures rebuilt the whole FastAPI app via create_app() on every test -- two mounted apps, every route and the OpenAPI schema, ~0.55s of per-test setup that dominated the api_fastapi suite. The app structure is identical for the default client, so build it once per session and reuse it. Because the app is now shared, per-test mutations are reset on every client fixture entry: restore the auth-manager singleton (the auth endpoint tests swap it for a mock without restoring) and clear any leftover dependency overrides. --- .../tests/unit/api_fastapi/conftest.py | 120 +++++++++++------- 1 file changed, 76 insertions(+), 44 deletions(-) diff --git a/airflow-core/tests/unit/api_fastapi/conftest.py b/airflow-core/tests/unit/api_fastapi/conftest.py index 03c43a178a090..d95d3d2fa9a57 100644 --- a/airflow-core/tests/unit/api_fastapi/conftest.py +++ b/airflow-core/tests/unit/api_fastapi/conftest.py @@ -23,9 +23,11 @@ import pytest import time_machine +from fastapi import FastAPI +from fastapi.routing import Mount from fastapi.testclient import TestClient -from airflow.api_fastapi.app import create_app +from airflow.api_fastapi.app import create_app, create_auth_manager from airflow.api_fastapi.auth.managers.simple.user import SimpleAuthManagerUser from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.models import Connection @@ -54,8 +56,16 @@ def get_api_path(request): return API_PATHS.get(subdirectory_name, "/") -@pytest.fixture -def test_client(request): +@pytest.fixture(scope="session") +def _shared_api_app(): + """ + Build the FastAPI app once per test session. + + ``create_app()`` rebuilds two full FastAPI apps (core + execution), registers every route and + builds the OpenAPI schema -- ~0.5s. The default ``test_client`` always uses the same config + (SimpleAuthManager), so the app structure is identical across tests; only per-test DB state and + request data differ. Building it once and reusing it removes that per-test rebuild cost. + """ with conf_vars( { ( @@ -64,54 +74,76 @@ def test_client(request): ): "airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager", } ): - app = create_app() - auth_manager: SimpleAuthManager = app.state.auth_manager - # set time_very_before to 2014-01-01 00:00:00 and time_very_after to tomorrow - # to make the JWT token always valid for all test cases with time_machine - time_very_before = datetime.datetime(2014, 1, 1, 0, 0, 0) - time_after = datetime.datetime.now() + datetime.timedelta(days=1) - with time_machine.travel(time_very_before, tick=False): - token = auth_manager._get_token_signer( - expiration_time_in_seconds=(time_after - time_very_before).total_seconds() - ).generate( - auth_manager.serialize_user( - SimpleAuthManagerUser(username="test", role="admin", teams=["team1"]) - ), - ) - with mock.patch("airflow.models.revoked_token.RevokedToken.is_revoked", return_value=False): - yield TestClient( - app, - headers={"Authorization": f"Bearer {token}"}, - base_url=f"{BASE_URL}{get_api_path(request)}", - ) + return create_app() + + +def _reset_shared_app(app: FastAPI) -> None: + """ + Reset mutable state on the session-shared app before each test. + + The app object is reused for the whole session, so any per-test mutation of its ``state`` or + ``dependency_overrides`` would leak into later tests. Two such mutations exist today and were + harmless only because each test used to build its own app: + + * the auth endpoint tests swap ``app.state.auth_manager`` for a mock without restoring it; + * several route tests (extra links, tasks, logs) install a ``dag_bag_from_app`` dependency + override in their per-test setup and never pop it. + + Restore the canonical auth manager and drop any leftover overrides (including on mounted + sub-apps) so each test starts from the pristine app and re-applies its own overrides. + """ + app.state.auth_manager = create_auth_manager() + app.dependency_overrides.clear() + for route in app.routes: + if isinstance(route, Mount) and isinstance(route.app, FastAPI): + route.app.dependency_overrides.clear() @pytest.fixture -def unauthenticated_test_client(request): - return TestClient(create_app(), base_url=f"{BASE_URL}{get_api_path(request)}") +def test_client(request, _shared_api_app): + app = _shared_api_app + _reset_shared_app(app) + auth_manager: SimpleAuthManager = app.state.auth_manager + # set time_very_before to 2014-01-01 00:00:00 and time_very_after to tomorrow + # to make the JWT token always valid for all test cases with time_machine + time_very_before = datetime.datetime(2014, 1, 1, 0, 0, 0) + time_after = datetime.datetime.now() + datetime.timedelta(days=1) + with time_machine.travel(time_very_before, tick=False): + token = auth_manager._get_token_signer( + expiration_time_in_seconds=(time_after - time_very_before).total_seconds() + ).generate( + auth_manager.serialize_user( + SimpleAuthManagerUser(username="test", role="admin", teams=["team1"]) + ), + ) + with mock.patch("airflow.models.revoked_token.RevokedToken.is_revoked", return_value=False): + yield TestClient( + app, + headers={"Authorization": f"Bearer {token}"}, + base_url=f"{BASE_URL}{get_api_path(request)}", + ) @pytest.fixture -def unauthorized_test_client(request): - with conf_vars( - { - ( - "core", - "auth_manager", - ): "airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager", - } - ): - app = create_app() - auth_manager: SimpleAuthManager = app.state.auth_manager - token = auth_manager._get_token_signer().generate( - auth_manager.serialize_user(SimpleAuthManagerUser(username="dummy", role=None)) +def unauthenticated_test_client(request, _shared_api_app): + _reset_shared_app(_shared_api_app) + return TestClient(_shared_api_app, base_url=f"{BASE_URL}{get_api_path(request)}") + + +@pytest.fixture +def unauthorized_test_client(request, _shared_api_app): + app = _shared_api_app + _reset_shared_app(app) + auth_manager: SimpleAuthManager = app.state.auth_manager + token = auth_manager._get_token_signer().generate( + auth_manager.serialize_user(SimpleAuthManagerUser(username="dummy", role=None)) + ) + with mock.patch("airflow.models.revoked_token.RevokedToken.is_revoked", return_value=False): + yield TestClient( + app, + headers={"Authorization": f"Bearer {token}"}, + base_url=f"{BASE_URL}{get_api_path(request)}", ) - with mock.patch("airflow.models.revoked_token.RevokedToken.is_revoked", return_value=False): - yield TestClient( - app, - headers={"Authorization": f"Bearer {token}"}, - base_url=f"{BASE_URL}{get_api_path(request)}", - ) @pytest.fixture From e2d2803538ce5dab437fdebd5b13cdac309b573b Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Tue, 9 Jun 2026 03:42:56 +0100 Subject: [PATCH 2/4] Add fresh_test_client for the DagBag-singleton test test_dagbag_used_as_singleton_in_dependency patches DBDagBag and asserts the constructor runs exactly once during app creation. With the session-shared app the app is built before the per-test patch, so the patched factory was never hit. Give that test a freshly built app (fresh_test_client) so it observes app construction. --- .../unit/api_fastapi/common/test_dagbag.py | 6 ++-- .../tests/unit/api_fastapi/conftest.py | 31 ++++++++++++++++--- 2 files changed, 30 insertions(+), 7 deletions(-) diff --git a/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py b/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py index cff9576beea4f..48c6f706ba7e2 100644 --- a/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py +++ b/airflow-core/tests/unit/api_fastapi/common/test_dagbag.py @@ -59,7 +59,7 @@ def factory(*args, **kwargs): purge_cached_app() yield - def test_dagbag_used_as_singleton_in_dependency(self, session, dag_maker, test_client): + def test_dagbag_used_as_singleton_in_dependency(self, session, dag_maker, fresh_test_client): """ Ensure DagBag is created only once and reused across multiple API requests. @@ -76,10 +76,10 @@ def test_dagbag_used_as_singleton_in_dependency(self, session, dag_maker, test_c BaseOperator(task_id="test_task") session.commit() - resp1 = test_client.get(f"/api/v2/dags/{dag_id}") + resp1 = fresh_test_client.get(f"/api/v2/dags/{dag_id}") assert resp1.status_code == 200 - resp2 = test_client.get(f"/api/v2/dags/{dag_id}") + resp2 = fresh_test_client.get(f"/api/v2/dags/{dag_id}") assert resp2.status_code == 200 assert self.dagbag_call_counter["count"] == 1 diff --git a/airflow-core/tests/unit/api_fastapi/conftest.py b/airflow-core/tests/unit/api_fastapi/conftest.py index d95d3d2fa9a57..4171bca6657fc 100644 --- a/airflow-core/tests/unit/api_fastapi/conftest.py +++ b/airflow-core/tests/unit/api_fastapi/conftest.py @@ -99,10 +99,7 @@ def _reset_shared_app(app: FastAPI) -> None: route.app.dependency_overrides.clear() -@pytest.fixture -def test_client(request, _shared_api_app): - app = _shared_api_app - _reset_shared_app(app) +def _authed_test_client(app: FastAPI, request): auth_manager: SimpleAuthManager = app.state.auth_manager # set time_very_before to 2014-01-01 00:00:00 and time_very_after to tomorrow # to make the JWT token always valid for all test cases with time_machine @@ -124,6 +121,32 @@ def test_client(request, _shared_api_app): ) +@pytest.fixture +def test_client(request, _shared_api_app): + _reset_shared_app(_shared_api_app) + yield from _authed_test_client(_shared_api_app, request) + + +@pytest.fixture +def fresh_test_client(request): + """ + Like ``test_client`` but backed by a freshly built app instead of the session-shared one. + + For the rare tests that patch app construction (e.g. counting ``DBDagBag`` instantiation) and + so need the app built *after* their patch is applied. + """ + with conf_vars( + { + ( + "core", + "auth_manager", + ): "airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager", + } + ): + app = create_app() + yield from _authed_test_client(app, request) + + @pytest.fixture def unauthenticated_test_client(request, _shared_api_app): _reset_shared_app(_shared_api_app) From 1d9e2ed9abccc87bc951ab888bfbfb4771a341d9 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Tue, 9 Jun 2026 12:31:56 +0100 Subject: [PATCH 3/4] Clear the shared DagBag cache between api_fastapi tests The session-shared app keeps app.state.dag_bag (an LRU/TTL cache keyed by dag_version_id) warm across the whole session. An entry warmed by one test would let a later test skip the serialized-Dag DB read, silently breaking query-count assertions (e.g. the grid ti_summaries stream tests) depending on execution order. Empty the cache on each client fixture entry, alongside the auth-manager and dependency-override resets. --- airflow-core/tests/unit/api_fastapi/conftest.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/airflow-core/tests/unit/api_fastapi/conftest.py b/airflow-core/tests/unit/api_fastapi/conftest.py index 4171bca6657fc..c9641a0565bdc 100644 --- a/airflow-core/tests/unit/api_fastapi/conftest.py +++ b/airflow-core/tests/unit/api_fastapi/conftest.py @@ -82,17 +82,22 @@ def _reset_shared_app(app: FastAPI) -> None: Reset mutable state on the session-shared app before each test. The app object is reused for the whole session, so any per-test mutation of its ``state`` or - ``dependency_overrides`` would leak into later tests. Two such mutations exist today and were + ``dependency_overrides`` would leak into later tests. These mutations exist today and were harmless only because each test used to build its own app: * the auth endpoint tests swap ``app.state.auth_manager`` for a mock without restoring it; * several route tests (extra links, tasks, logs) install a ``dag_bag_from_app`` dependency - override in their per-test setup and never pop it. - - Restore the canonical auth manager and drop any leftover overrides (including on mounted - sub-apps) so each test starts from the pristine app and re-applies its own overrides. + override in their per-test setup and never pop it; + * ``app.state.dag_bag`` carries an LRU/TTL cache of deserialized Dags keyed by + ``dag_version_id``. A warm entry from an earlier test would let a later test skip the + serialized-Dag DB read, which silently breaks query-count assertions (e.g. the grid + ``ti_summaries`` stream tests) depending on execution order. + + Restore the canonical auth manager, drop leftover overrides (including on mounted sub-apps), + and empty the Dag cache so each test starts from the pristine app and re-applies its own state. """ app.state.auth_manager = create_auth_manager() + app.state.dag_bag.clear_cache() app.dependency_overrides.clear() for route in app.routes: if isinstance(route, Mount) and isinstance(route.app, FastAPI): From 50dc240da8a5d527b852c187523f6790e2529b5b Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Tue, 9 Jun 2026 20:21:30 +0100 Subject: [PATCH 4/4] Snapshot/restore shared app state instead of enumerated resets Per review: replace _reset_shared_app with an _isolated_shared_app yield fixture that snapshots state + dependency_overrides on the root app and every mounted sub-app (recursively) on entry and restores them on teardown, so isolation is resilient to future mutations rather than enumerating the known ones. app.state.dag_bag is the exception: its cache is mutated in place (not a state rebinding), so a snapshot can't empty it -- it is still cleared explicitly. --- .../tests/unit/api_fastapi/conftest.py | 77 +++++++++++-------- 1 file changed, 44 insertions(+), 33 deletions(-) diff --git a/airflow-core/tests/unit/api_fastapi/conftest.py b/airflow-core/tests/unit/api_fastapi/conftest.py index c9641a0565bdc..f275d48f967be 100644 --- a/airflow-core/tests/unit/api_fastapi/conftest.py +++ b/airflow-core/tests/unit/api_fastapi/conftest.py @@ -27,7 +27,7 @@ from fastapi.routing import Mount from fastapi.testclient import TestClient -from airflow.api_fastapi.app import create_app, create_auth_manager +from airflow.api_fastapi.app import create_app from airflow.api_fastapi.auth.managers.simple.user import SimpleAuthManagerUser from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.models import Connection @@ -77,31 +77,45 @@ def _shared_api_app(): return create_app() -def _reset_shared_app(app: FastAPI) -> None: - """ - Reset mutable state on the session-shared app before each test. - - The app object is reused for the whole session, so any per-test mutation of its ``state`` or - ``dependency_overrides`` would leak into later tests. These mutations exist today and were - harmless only because each test used to build its own app: - - * the auth endpoint tests swap ``app.state.auth_manager`` for a mock without restoring it; - * several route tests (extra links, tasks, logs) install a ``dag_bag_from_app`` dependency - override in their per-test setup and never pop it; - * ``app.state.dag_bag`` carries an LRU/TTL cache of deserialized Dags keyed by - ``dag_version_id``. A warm entry from an earlier test would let a later test skip the - serialized-Dag DB read, which silently breaks query-count assertions (e.g. the grid - ``ti_summaries`` stream tests) depending on execution order. - - Restore the canonical auth manager, drop leftover overrides (including on mounted sub-apps), - and empty the Dag cache so each test starts from the pristine app and re-applies its own state. - """ - app.state.auth_manager = create_auth_manager() - app.state.dag_bag.clear_cache() - app.dependency_overrides.clear() +def _mounted_fastapi_apps(app: FastAPI) -> list[FastAPI]: + """Return ``app`` and every FastAPI app mounted under it, recursively (``/execution``, ``/auth``, ...).""" + apps = [app] for route in app.routes: if isinstance(route, Mount) and isinstance(route.app, FastAPI): - route.app.dependency_overrides.clear() + apps.extend(_mounted_fastapi_apps(route.app)) + return apps + + +@pytest.fixture +def _isolated_shared_app(_shared_api_app): + """ + Yield the session-shared app with its mutable state snapshotted and restored around each test. + + The app is built once per session, so a test that rebinds something on ``app.state`` (the auth + endpoint tests swap ``auth_manager`` for a mock) or installs a ``dependency_overrides`` entry + (extra-links/tasks/logs install a ``dag_bag_from_app`` override) would leak into later tests. + Snapshotting ``state`` and ``dependency_overrides`` on the root app and every mounted sub-app on + entry and restoring them on exit keeps the reset resilient to future mutations without having to + enumerate them. + + ``app.state.dag_bag`` is the exception: tests mutate the DagBag object *in place* (its cache of + deserialized Dags fills as requests resolve them), which a state snapshot can't undo, so its + cache is cleared explicitly. A leaked warm entry would otherwise let a later test skip a + serialized-Dag DB read and break query-count assertions (e.g. the grid ``ti_summaries`` stream + tests) depending on execution order. + """ + apps = _mounted_fastapi_apps(_shared_api_app) + # ``app.state._state`` is Starlette's backing dict for ``State`` -- the only way to enumerate it. + saved = [(app, dict(app.state._state), dict(app.dependency_overrides)) for app in apps] + _shared_api_app.state.dag_bag.clear_cache() + try: + yield _shared_api_app + finally: + for app, state, overrides in saved: + app.state._state.clear() + app.state._state.update(state) + app.dependency_overrides.clear() + app.dependency_overrides.update(overrides) def _authed_test_client(app: FastAPI, request): @@ -127,9 +141,8 @@ def _authed_test_client(app: FastAPI, request): @pytest.fixture -def test_client(request, _shared_api_app): - _reset_shared_app(_shared_api_app) - yield from _authed_test_client(_shared_api_app, request) +def test_client(request, _isolated_shared_app): + yield from _authed_test_client(_isolated_shared_app, request) @pytest.fixture @@ -153,15 +166,13 @@ def fresh_test_client(request): @pytest.fixture -def unauthenticated_test_client(request, _shared_api_app): - _reset_shared_app(_shared_api_app) - return TestClient(_shared_api_app, base_url=f"{BASE_URL}{get_api_path(request)}") +def unauthenticated_test_client(request, _isolated_shared_app): + return TestClient(_isolated_shared_app, base_url=f"{BASE_URL}{get_api_path(request)}") @pytest.fixture -def unauthorized_test_client(request, _shared_api_app): - app = _shared_api_app - _reset_shared_app(app) +def unauthorized_test_client(request, _isolated_shared_app): + app = _isolated_shared_app auth_manager: SimpleAuthManager = app.state.auth_manager token = auth_manager._get_token_signer().generate( auth_manager.serialize_user(SimpleAuthManagerUser(username="dummy", role=None))