This is an automated email from the ASF dual-hosted git repository.
dheerajturaga pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6445ba58062 Add tests for edge3 worker API app and JWT auth (#72482)
6445ba58062 is described below
commit 6445ba58062fe03c3e0d2d111d4784b475f5ad9d
Author: Dheeren Mohta <[email protected]>
AuthorDate: Sun Sep 6 09:09:05 2026 +0530
Add tests for edge3 worker API app and JWT auth (#72482)
The FastAPI app factory and the JWT authorization layer for the Edge
Worker API had no dedicated unit tests, leaving the route/mount wiring
(v1 API prefix, UI prefix, static/res mounts, .cjs mimetype) and the
authorization flow (method-claim checks, per-exception anonymized 403
responses, validator caching semantics) unverified against regressions.
Closes: #72266
---
.../tests/unit/always/test_project_structure.py | 2 -
.../edge3/tests/unit/edge3/worker_api/test_app.py | 65 +++++++++
.../edge3/tests/unit/edge3/worker_api/test_auth.py | 152 +++++++++++++++++++++
3 files changed, 217 insertions(+), 2 deletions(-)
diff --git a/airflow-core/tests/unit/always/test_project_structure.py
b/airflow-core/tests/unit/always/test_project_structure.py
index bc1bca06407..fe0d230deac 100644
--- a/airflow-core/tests/unit/always/test_project_structure.py
+++ b/airflow-core/tests/unit/always/test_project_structure.py
@@ -109,8 +109,6 @@ class TestProjectStructure:
"providers/edge3/tests/unit/edge3/models/test_edge_job.py",
"providers/edge3/tests/unit/edge3/models/test_edge_logs.py",
"providers/edge3/tests/unit/edge3/models/test_edge_worker.py",
- "providers/edge3/tests/unit/edge3/worker_api/test_app.py",
- "providers/edge3/tests/unit/edge3/worker_api/test_auth.py",
"providers/edge3/tests/unit/edge3/worker_api/test_datamodels.py",
"providers/edge3/tests/unit/edge3/worker_api/test_datamodels_ui.py",
"providers/fab/tests/unit/fab/auth_manager/api_fastapi/datamodels/test_login.py",
diff --git a/providers/edge3/tests/unit/edge3/worker_api/test_app.py
b/providers/edge3/tests/unit/edge3/worker_api/test_app.py
new file mode 100644
index 00000000000..ec1c54522c9
--- /dev/null
+++ b/providers/edge3/tests/unit/edge3/worker_api/test_app.py
@@ -0,0 +1,65 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import mimetypes
+from unittest import mock
+
+import pytest
+from fastapi.routing import APIRoute
+from starlette.routing import Mount
+
+from airflow.providers.edge3.worker_api.app import create_edge_worker_api_app
+
+
[email protected]
+def app():
+ # The React UI bundle (``plugins/www/dist``) is only produced by a JS
build step and
+ # is not present in a plain source checkout, so the real StaticFiles
directory check
+ # would fail here regardless of how create_edge_worker_api_app() itself
behaves.
+ with mock.patch("airflow.providers.edge3.worker_api.app.StaticFiles") as
mocked_static_files:
+ mocked_static_files.side_effect = lambda *args, **kwargs:
mock.MagicMock()
+ yield create_edge_worker_api_app()
+
+
+def _api_route_paths(app) -> set[str]:
+ return {route.path for route in app.routes if isinstance(route, APIRoute)}
+
+
+class TestCreateEdgeWorkerApiApp:
+ def test_v1_routers_are_mounted_under_the_v1_prefix(self, app):
+ route_paths = _api_route_paths(app)
+ assert "/v1/health" in route_paths
+ assert "/v1/jobs/fetch/{worker_name}" in route_paths
+ assert
"/v1/logs/logfile_path/{dag_id}/{task_id}/{run_id}/{try_number}/{map_index}" in
route_paths
+ assert "/v1/worker/{worker_name}" in route_paths
+
+ def test_ui_router_is_mounted_under_the_ui_prefix(self, app):
+ route_paths = _api_route_paths(app)
+ assert "/ui/worker" in route_paths
+ assert not any(path.startswith("/v1") for path in route_paths if path
== "/ui/worker")
+
+ def test_static_and_res_directories_are_mounted(self, app):
+ mounts = {route.path: route for route in app.routes if
isinstance(route, Mount)}
+ assert "/static" in mounts
+ assert "/res" in mounts
+ assert mounts["/static"].name == "react_static_plugin_files"
+ assert mounts["/res"].name == "react_res_plugin_files"
+
+ def test_cjs_mimetype_is_registered_as_javascript(self, app):
+ # Serving .cjs with the wrong mimetype breaks the Edge Worker UI in
the browser.
+ assert mimetypes.guess_type("plugin.cjs")[0] ==
"application/javascript"
diff --git a/providers/edge3/tests/unit/edge3/worker_api/test_auth.py
b/providers/edge3/tests/unit/edge3/worker_api/test_auth.py
new file mode 100644
index 00000000000..3d4263be855
--- /dev/null
+++ b/providers/edge3/tests/unit/edge3/worker_api/test_auth.py
@@ -0,0 +1,152 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest import mock
+
+import jwt
+import pytest
+from fastapi import HTTPException, Request
+from itsdangerous import BadSignature
+
+from airflow.api_fastapi.auth.tokens import JWTGenerator
+from airflow.providers.edge3.worker_api.auth import (
+ jwt_token_authorization,
+ jwt_token_authorization_rest,
+ jwt_validator,
+)
+
+from tests_common.test_utils.config import conf_vars
+
+JWT_SECRET = "test-jwt-secret"
+
+
+def _token(method: str | None = "test.method", secret: str = JWT_SECRET) ->
str:
+ # Mirrors how
providers/edge3/src/airflow/providers/edge3/cli/api_client.py generates
+ # tokens for real edge workers, so the token shape (iss/aud/exp/...
claims) matches
+ # what jwt_token_authorization() actually has to validate in production.
+ generator = JWTGenerator(secret_key=secret, valid_for=300, audience="api")
+ return generator.generate(extras={"method": method} if method is not None
else {})
+
+
[email protected](autouse=True)
+def _reset_jwt_validator_cache():
+ # jwt_validator() is cached: make sure config overrides in one test can
never leak
+ # a stale validator into another.
+ jwt_validator.cache_clear()
+ yield
+ jwt_validator.cache_clear()
+
+
+class TestJwtTokenAuthorization:
+ @conf_vars({("api_auth", "jwt_secret"): JWT_SECRET, ("api_auth",
"jwt_leeway"): "5"})
+ def test_matching_method_claim_is_authorized(self):
+ jwt_token_authorization("test.method", _token("test.method"))
+
+ @conf_vars({("api_auth", "jwt_secret"): JWT_SECRET, ("api_auth",
"jwt_leeway"): "5"})
+ def test_missing_method_claim_is_forbidden(self):
+ with pytest.raises(HTTPException) as exc_info:
+ jwt_token_authorization("test.method", _token(method=None))
+ assert exc_info.value.status_code == 403
+
+ @conf_vars({("api_auth", "jwt_secret"): JWT_SECRET, ("api_auth",
"jwt_leeway"): "5"})
+ def test_mismatched_method_claim_is_forbidden(self):
+ with pytest.raises(HTTPException) as exc_info:
+ jwt_token_authorization("test.method", _token("other.method"))
+ assert exc_info.value.status_code == 403
+
+
+class TestJwtTokenAuthorizationForbiddenResponse:
+ """
+ Every handled failure is collapsed onto the same anonymised 403 response by
+ ``_forbidden_response`` so callers can never distinguish *why* a token was
rejected -
+ only the server-side log carries the real reason.
+ """
+
+ @pytest.mark.parametrize(
+ "error",
+ [
+ BadSignature("Signature does not match"),
+ jwt.InvalidAudienceError("Invalid audience"),
+ jwt.InvalidSignatureError("Signature verification failed"),
+ jwt.ImmatureSignatureError("The token is not yet valid"),
+ jwt.ExpiredSignatureError("Signature has expired"),
+ jwt.InvalidIssuedAtError("Issued at claim is in the future"),
+ ValueError("Some other unexpected failure"),
+ ],
+ )
+ @mock.patch("airflow.providers.edge3.worker_api.auth.jwt_validate",
autospec=True)
+ def test_each_handled_failure_is_forbidden_and_anonymized(self,
mock_jwt_validate, error):
+ mock_jwt_validate.side_effect = error
+
+ with pytest.raises(HTTPException) as exc_info:
+ jwt_token_authorization("test.method", "some-token")
+
+ assert exc_info.value.status_code == 403
+ assert "error_id=" in exc_info.value.detail
+ assert str(error) not in exc_info.value.detail
+
+
+class TestJwtValidatorCaching:
+ def test_validator_is_cached_and_reuses_previously_configured_secret(self):
+ with conf_vars({("api_auth", "jwt_secret"): "secret-one"}):
+ first = jwt_validator()
+
+ with conf_vars({("api_auth", "jwt_secret"): "secret-two"}):
+ # No cache_clear() call: the cached validator (built with
"secret-one")
+ # is silently reused, per the caveat called out for this function.
+ second = jwt_validator()
+
+ assert first is second
+ assert second.secret_key == "secret-one"
+
+ def test_cache_clear_picks_up_the_new_secret(self):
+ with conf_vars({("api_auth", "jwt_secret"): "secret-one"}):
+ first = jwt_validator()
+
+ jwt_validator.cache_clear()
+
+ with conf_vars({("api_auth", "jwt_secret"): "secret-two"}):
+ second = jwt_validator()
+
+ assert second is not first
+ assert second.secret_key == "secret-two"
+
+ @conf_vars({("api_auth", "jwt_secret"): JWT_SECRET, ("api_auth",
"jwt_leeway"): "90"})
+ def test_leeway_is_read_from_config(self):
+ assert jwt_validator().leeway == 90
+
+
+class TestJwtTokenAuthorizationRest:
+ @pytest.mark.parametrize(
+ ("path", "expected_method"),
+ [
+ ("/edge_worker/v1/jobs/fetch/worker1", "jobs/fetch/worker1"),
+ ("/edge_worker/v1/health", "health"),
+ ("/some/other/path", "/some/other/path"),
+ ],
+ )
+
@mock.patch("airflow.providers.edge3.worker_api.auth.jwt_token_authorization",
autospec=True)
+ def test_strips_edge_worker_v1_prefix_and_falls_back_to_full_path(
+ self, mock_jwt_token_authorization, path, expected_method
+ ):
+ request = mock.MagicMock(spec=Request)
+ request.url.path = path
+
+ jwt_token_authorization_rest(request, authorization="some-token")
+
+ mock_jwt_token_authorization.assert_called_once_with(expected_method,
"some-token")
