GayathriSrividya commented on code in PR #68499: URL: https://github.com/apache/airflow/pull/68499#discussion_r4053620569
########## airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py: ########## @@ -20,17 +20,76 @@ from unittest.mock import AsyncMock import pytest -from fastapi import FastAPI +import svcs +from fastapi import FastAPI, HTTPException, Request, status +from fastapi.security import HTTPBearer +from fastapi.testclient import TestClient from airflow.api_fastapi.auth.tokens import JWTValidator from airflow.api_fastapi.execution_api.app import lifespan +from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken +from airflow.api_fastapi.execution_api.security import ( + _REQUEST_SCOPE_TOKEN_KEY, + _jwt_bearer, +) + + [email protected] +def jwt_bearer_client(): + """Test client that exercises JWTBearer so request.scope is populated for the middleware.""" + from starlette.routing import Mount + + from airflow.api_fastapi.app import cached_app + + app = cached_app(apps="execution") + + exec_app: FastAPI | None = None + for route in app.routes: + if isinstance(route, Mount) and route.path == "/execution" and isinstance(route.app, FastAPI): + exec_app = route.app + break + if exec_app is None: + raise RuntimeError("Execution API sub-app not found") + + _http_bearer = HTTPBearer(auto_error=False) + + async def mock_jwt_bearer(request: Request): + """Drop-in for _jwt_bearer that uses the registered JWTValidator mock and sets scope.""" + if cached := request.scope.get(_REQUEST_SCOPE_TOKEN_KEY): + return cached + + creds = await _http_bearer(request) + if not creds: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing auth token") + + async with svcs.Container(request.app.state.svcs_registry) as services: + validator: JWTValidator = await services.aget(JWTValidator) + try: + claims = await validator.avalidated_claims(creds.credentials, {}) + except Exception: + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Invalid auth token") + + claims.setdefault("scope", "execution") + token = TIToken(id=claims["sub"], claims=TIClaims(**claims)) + request.scope[_REQUEST_SCOPE_TOKEN_KEY] = token + return token + + exec_app.dependency_overrides[_jwt_bearer] = mock_jwt_bearer + + with TestClient(app) as c: + yield c + + exec_app.dependency_overrides.pop(_jwt_bearer, None) @pytest.fixture -def exec_app(client): - last_route = client.app.routes[-1] - assert isinstance(last_route.app, FastAPI) - return last_route.app +def exec_app(jwt_bearer_client): + from starlette.routing import Mount Review Comment: moved the inline imports to the module level in the test file as requested. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
