From 31b07eb4e33d359e5aef2523eeeb1f2843cf791f Mon Sep 17 00:00:00 2001 From: vincbeck Date: Thu, 30 Jan 2025 11:38:12 -0500 Subject: [PATCH] Update AWS auth manager to use Fastapi instead of Flask --- airflow/config_templates/config.yml | 12 ++ .../aws/auth_manager/aws_auth_manager.py | 53 +++--- .../{views => router}/__init__.py | 0 .../amazon/aws/auth_manager/router/login.py | 124 ++++++++++++++ .../amazon/aws/auth_manager/views/auth.py | 151 ------------------ .../{views => router}/__init__.py | 0 .../test_auth.py => router/test_login.py} | 104 +++++------- .../aws/auth_manager/test_aws_auth_manager.py | 58 +------ 8 files changed, 214 insertions(+), 288 deletions(-) rename providers/src/airflow/providers/amazon/aws/auth_manager/{views => router}/__init__.py (100%) create mode 100644 providers/src/airflow/providers/amazon/aws/auth_manager/router/login.py delete mode 100644 providers/src/airflow/providers/amazon/aws/auth_manager/views/auth.py rename providers/tests/amazon/aws/auth_manager/{views => router}/__init__.py (100%) rename providers/tests/amazon/aws/auth_manager/{views/test_auth.py => router/test_login.py} (57%) diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml index 918b111acd43a..a46bad1c44bc6 100644 --- a/airflow/config_templates/config.yml +++ b/airflow/config_templates/config.yml @@ -2676,3 +2676,15 @@ dag_processor: type: integer example: ~ default: "30" +fastapi: + description: Configuration for the Fastapi webserver. + options: + base_url: + description: | + The base url of the Fastapi endpoint. Airflow cannot guess what domain or CNAME you are using. + If the Airflow console (the front-end) and the Fastapi apis are on a different domain, this config + should contain the Fastapi apis endpoint. + version_added: ~ + type: string + example: ~ + default: "http://localhost:29091" diff --git a/providers/src/airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py b/providers/src/airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py index b93d6a1ad3789..6561e8121c465 100644 --- a/providers/src/airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py +++ b/providers/src/airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py @@ -20,9 +20,10 @@ from collections import defaultdict from collections.abc import Container, Sequence from functools import cached_property -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, cast -from flask import session, url_for +from fastapi import FastAPI +from flask import session from airflow.auth.managers.base_auth_manager import BaseAuthManager from airflow.auth.managers.models.resource_details import ( @@ -34,6 +35,7 @@ VariableDetails, ) from airflow.cli.cli_config import CLICommand, DefaultHelpParser, GroupCommand +from airflow.configuration import conf from airflow.exceptions import AirflowOptionalProviderFeatureException from airflow.providers.amazon.aws.auth_manager.avp.entities import AvpEntities from airflow.providers.amazon.aws.auth_manager.avp.facade import ( @@ -43,11 +45,7 @@ from airflow.providers.amazon.aws.auth_manager.cli.definition import ( AWS_AUTH_MANAGER_COMMANDS, ) -from airflow.providers.amazon.aws.auth_manager.security_manager.aws_security_manager_override import ( - AwsSecurityManagerOverride, -) from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser -from airflow.providers.amazon.aws.auth_manager.views.auth import AwsAuthManagerAuthenticationViews from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS if TYPE_CHECKING: @@ -61,7 +59,6 @@ IsAuthorizedVariableRequest, ) from airflow.auth.managers.models.resource_details import AssetDetails, ConfigurationDetails - from airflow.www.extensions.init_appbuilder import AirflowAppBuilder class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]): @@ -72,8 +69,6 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]): authentication and authorization in Airflow. """ - appbuilder: AirflowAppBuilder | None = None - def __init__(self) -> None: if not AIRFLOW_V_3_0_PLUS: raise AirflowOptionalProviderFeatureException( @@ -87,12 +82,27 @@ def __init__(self) -> None: def avp_facade(self): return AwsAuthManagerAmazonVerifiedPermissionsFacade() + @cached_property + def fastapi_endpoint(self) -> str: + return conf.get("fastapi", "base_url") + def get_user(self) -> AwsAuthManagerUser | None: return session["aws_user"] if self.is_logged_in() else None def is_logged_in(self) -> bool: return "aws_user" in session + def deserialize_user(self, token: dict[str, Any]) -> AwsAuthManagerUser: + return AwsAuthManagerUser(**token) + + def serialize_user(self, user: AwsAuthManagerUser) -> dict[str, Any]: + return { + "user_id": user.get_id(), + "groups": user.get_groups(), + "username": user.username, + "email": user.email, + } + def is_authorized_configuration( self, *, @@ -367,14 +377,10 @@ def _has_access_to_menu_item(request: IsAuthorizedRequest): return accessible_items def get_url_login(self, **kwargs) -> str: - return url_for("AwsAuthManagerAuthenticationViews.login") + return f"{self.fastapi_endpoint}/auth/login" def get_url_logout(self) -> str: - return url_for("AwsAuthManagerAuthenticationViews.logout") - - @cached_property - def security_manager(self) -> AwsSecurityManagerOverride: - return AwsSecurityManagerOverride(self.appbuilder) + raise NotImplementedError() @staticmethod def get_cli_commands() -> list[CLICommand]: @@ -387,9 +393,20 @@ def get_cli_commands() -> list[CLICommand]: ), ] - def register_views(self) -> None: - if self.appbuilder: - self.appbuilder.add_view_no_menu(AwsAuthManagerAuthenticationViews()) + def get_fastapi_app(self) -> FastAPI | None: + from airflow.providers.amazon.aws.auth_manager.router.login import login_router + + app = FastAPI( + title="AWS auth manager sub application", + description=( + "This is the AWS auth manager fastapi sub application. This API is only available if the " + "auth manager used in the Airflow environment is AWS auth manager. " + "This sub application provides login routes." + ), + ) + app.include_router(login_router) + + return app @staticmethod def _get_menu_item_request(resource_name: str) -> IsAuthorizedRequest: diff --git a/providers/src/airflow/providers/amazon/aws/auth_manager/views/__init__.py b/providers/src/airflow/providers/amazon/aws/auth_manager/router/__init__.py similarity index 100% rename from providers/src/airflow/providers/amazon/aws/auth_manager/views/__init__.py rename to providers/src/airflow/providers/amazon/aws/auth_manager/router/__init__.py diff --git a/providers/src/airflow/providers/amazon/aws/auth_manager/router/login.py b/providers/src/airflow/providers/amazon/aws/auth_manager/router/login.py new file mode 100644 index 0000000000000..aca770dca29e6 --- /dev/null +++ b/providers/src/airflow/providers/amazon/aws/auth_manager/router/login.py @@ -0,0 +1,124 @@ +# 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 logging +from typing import Any + +import anyio +from fastapi import HTTPException, Request +from starlette import status +from starlette.responses import RedirectResponse + +from airflow.api_fastapi.app import get_auth_manager +from airflow.api_fastapi.common.router import AirflowRouter +from airflow.configuration import conf +from airflow.providers.amazon.aws.auth_manager.constants import CONF_SAML_METADATA_URL_KEY, CONF_SECTION_NAME +from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser + +try: + from onelogin.saml2.auth import OneLogin_Saml2_Auth + from onelogin.saml2.errors import OneLogin_Saml2_Error + from onelogin.saml2.idp_metadata_parser import OneLogin_Saml2_IdPMetadataParser +except ImportError: + raise ImportError( + "AWS auth manager requires the python3-saml library but it is not installed by default. " + "Please install the python3-saml library by running: " + "pip install apache-airflow-providers-amazon[python3-saml]" + ) + +log = logging.getLogger(__name__) +login_router = AirflowRouter(tags=["AWSAuthManagerLogin"]) + + +@login_router.get("/login") +def login(request: Request): + """Authenticate the user.""" + saml_auth = _init_saml_auth(request) + callback_url = saml_auth.login() + return RedirectResponse(url=callback_url) + + +@login_router.post("/login_callback") +def login_callback(request: Request): + """Authenticate the user.""" + saml_auth = _init_saml_auth(request) + try: + saml_auth.process_response() + except OneLogin_Saml2_Error as e: + log.exception(e) + raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, "Failed to authenticate") + errors = saml_auth.get_errors() + is_authenticated = saml_auth.is_authenticated() + if not is_authenticated: + error_reason = saml_auth.get_last_error_reason() + log.error("Failed to authenticate") + log.error("Errors: %s", errors) + log.error("Error reason: %s", error_reason) + raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, f"Failed to authenticate: {error_reason}") + + attributes = saml_auth.get_attributes() + user = AwsAuthManagerUser( + user_id=attributes["id"][0], + groups=attributes["groups"], + username=saml_auth.get_nameid(), + email=attributes["email"][0] if "email" in attributes else None, + ) + return RedirectResponse(url=f"/webapp?token={get_auth_manager().get_jwt_token(user)}", status_code=303) + + +def _init_saml_auth(request: Request) -> OneLogin_Saml2_Auth: + request_data = _prepare_request(request) + base_url = conf.get(section="fastapi", key="base_url") + settings = { + # We want to keep this flag on in case of errors. + # It provides an error reasons, if turned off, it does not + "debug": True, + "sp": { + "entityId": "aws-auth-manager-saml-client", + "assertionConsumerService": { + "url": f"{base_url}/auth/login_callback", + "binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST", + }, + }, + } + merged_settings = OneLogin_Saml2_IdPMetadataParser.merge_settings(_get_idp_data(), settings) + return OneLogin_Saml2_Auth(request_data, merged_settings) + + +def _prepare_request(request: Request) -> dict: + host = request.headers.get("host", request.client.host if request.client else "localhost") + data: dict[str, Any] = { + "https": "on" if request.url.scheme == "https" else "off", + "http_host": host, + "server_port": request.url.port, + "script_name": request.url.path, + "get_data": request.query_params, + "post_data": {}, + } + form_data = anyio.from_thread.run(request.form) + if "SAMLResponse" in form_data: + data["post_data"]["SAMLResponse"] = form_data["SAMLResponse"] + if "RelayState" in form_data: + data["post_data"]["RelayState"] = form_data["RelayState"] + return data + + +def _get_idp_data() -> dict: + saml_metadata_url = conf.get_mandatory_value(CONF_SECTION_NAME, CONF_SAML_METADATA_URL_KEY) + return OneLogin_Saml2_IdPMetadataParser.parse_remote(saml_metadata_url) diff --git a/providers/src/airflow/providers/amazon/aws/auth_manager/views/auth.py b/providers/src/airflow/providers/amazon/aws/auth_manager/views/auth.py deleted file mode 100644 index e08c2a7a6e100..0000000000000 --- a/providers/src/airflow/providers/amazon/aws/auth_manager/views/auth.py +++ /dev/null @@ -1,151 +0,0 @@ -# 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 logging -from functools import cached_property - -from flask import make_response, redirect, request, session, url_for -from flask_appbuilder import expose - -from airflow.configuration import conf -from airflow.exceptions import AirflowException -from airflow.providers.amazon.aws.auth_manager.constants import CONF_SAML_METADATA_URL_KEY, CONF_SECTION_NAME -from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser -from airflow.www.app import csrf -from airflow.www.views import AirflowBaseView - -try: - from onelogin.saml2.auth import OneLogin_Saml2_Auth - from onelogin.saml2.idp_metadata_parser import OneLogin_Saml2_IdPMetadataParser -except ImportError: - raise ImportError( - "AWS auth manager requires the python3-saml library but it is not installed by default. " - "Please install the python3-saml library by running: " - "pip install apache-airflow-providers-amazon[python3-saml]" - ) - -logger = logging.getLogger(__name__) - - -class AwsAuthManagerAuthenticationViews(AirflowBaseView): - """ - Views specific to AWS auth manager authentication mechanism. - - Some code below is inspired from - https://github.com/SAML-Toolkits/python3-saml/blob/6988bdab7a203abfe8dc264992f7e350c67aef3d/demo-flask/index.py - """ - - @cached_property - def idp_data(self) -> dict: - saml_metadata_url = conf.get_mandatory_value(CONF_SECTION_NAME, CONF_SAML_METADATA_URL_KEY) - return OneLogin_Saml2_IdPMetadataParser.parse_remote(saml_metadata_url) - - @expose("/login") - def login(self): - """Start login process.""" - saml_auth = self._init_saml_auth() - return redirect(saml_auth.login()) - - @expose("/logout", methods=("GET", "POST")) - def logout(self): - """Start logout process.""" - session.clear() - saml_auth = self._init_saml_auth() - - return redirect(saml_auth.logout()) - - @csrf.exempt - @expose("/login_callback", methods=("GET", "POST")) - def login_callback(self): - """ - Redirect the user to this callback after successful login. - - CSRF protection needs to be disabled otherwise the callback won't work. - """ - saml_auth = self._init_saml_auth() - saml_auth.process_response() - errors = saml_auth.get_errors() - is_authenticated = saml_auth.is_authenticated() - if not is_authenticated: - error_reason = saml_auth.get_last_error_reason() - logger.error("Failed to authenticate") - logger.error("Errors: %s", errors) - logger.error("Error reason: %s", error_reason) - raise AirflowException(f"Failed to authenticate: {error_reason}") - - attributes = saml_auth.get_attributes() - user = AwsAuthManagerUser( - user_id=attributes["id"][0], - groups=attributes["groups"], - username=saml_auth.get_nameid(), - email=attributes["email"][0] if "email" in attributes else None, - ) - session["aws_user"] = user - - return redirect(url_for("Airflow.index")) - - @csrf.exempt - @expose("/logout_callback", methods=("GET", "POST")) - def logout_callback(self): - raise NotImplementedError("AWS Identity center does not support SLO (Single Logout Service)") - - @expose("/login_metadata") - def login_metadata(self): - saml_auth = self._init_saml_auth() - settings = saml_auth.get_settings() - metadata = settings.get_sp_metadata() - errors = settings.validate_metadata(metadata) - - if len(errors) == 0: - resp = make_response(metadata, 200) - resp.headers["Content-Type"] = "text/xml" - else: - resp = make_response(", ".join(errors), 500) - return resp - - @staticmethod - def _prepare_flask_request() -> dict: - return { - "https": "on" if request.scheme == "https" else "off", - "http_host": request.host, - "script_name": request.path, - "get_data": request.args.copy(), - "post_data": request.form.copy(), - } - - def _init_saml_auth(self) -> OneLogin_Saml2_Auth: - request_data = self._prepare_flask_request() - base_url = conf.get(section="webserver", key="base_url") - settings = { - # We want to keep this flag on in case of errors. - # It provides an error reasons, if turned off, it does not - "debug": True, - "sp": { - "entityId": f"{base_url}/login_metadata", - "assertionConsumerService": { - "url": f"{base_url}/login_callback", - "binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST", - }, - "singleLogoutService": { - "url": f"{base_url}/logout_callback", - "binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect", - }, - }, - } - merged_settings = OneLogin_Saml2_IdPMetadataParser.merge_settings(settings, self.idp_data) - return OneLogin_Saml2_Auth(request_data, merged_settings) diff --git a/providers/tests/amazon/aws/auth_manager/views/__init__.py b/providers/tests/amazon/aws/auth_manager/router/__init__.py similarity index 100% rename from providers/tests/amazon/aws/auth_manager/views/__init__.py rename to providers/tests/amazon/aws/auth_manager/router/__init__.py diff --git a/providers/tests/amazon/aws/auth_manager/views/test_auth.py b/providers/tests/amazon/aws/auth_manager/router/test_login.py similarity index 57% rename from providers/tests/amazon/aws/auth_manager/views/test_auth.py rename to providers/tests/amazon/aws/auth_manager/router/test_login.py index 2521dd9e43e6c..f68eb7b7fe23e 100644 --- a/providers/tests/amazon/aws/auth_manager/views/test_auth.py +++ b/providers/tests/amazon/aws/auth_manager/router/test_login.py @@ -19,11 +19,16 @@ from unittest.mock import Mock, patch import pytest -from flask import session, url_for -from airflow.exceptions import AirflowException from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS -from airflow.www import app as application + +if not AIRFLOW_V_3_0_PLUS: + pytest.skip("AWS auth manager is only compatible with Airflow >= 3.0.0", allow_module_level=True) + +from fastapi.testclient import TestClient +from onelogin.saml2.idp_metadata_parser import OneLogin_Saml2_IdPMetadataParser + +from airflow.api_fastapi.app import create_app from tests_common.test_utils.config import conf_vars @@ -47,7 +52,7 @@ @pytest.fixture -def aws_app(): +def test_client(): with conf_vars( { ( @@ -58,42 +63,26 @@ def aws_app(): } ): with ( - patch( - "airflow.providers.amazon.aws.auth_manager.views.auth.OneLogin_Saml2_IdPMetadataParser" - ) as mock_parser, + patch.object(OneLogin_Saml2_IdPMetadataParser, "parse_remote") as mock_parse_remote, patch( "airflow.providers.amazon.aws.auth_manager.avp.facade.AwsAuthManagerAmazonVerifiedPermissionsFacade.is_policy_store_schema_up_to_date" ) as mock_is_policy_store_schema_up_to_date, ): mock_is_policy_store_schema_up_to_date.return_value = True - mock_parser.parse_remote.return_value = SAML_METADATA_PARSED - return application.create_app(testing=True, config={"WTF_CSRF_ENABLED": False}) - - -@pytest.mark.skipif( - not AIRFLOW_V_3_0_PLUS, reason="AWS auth manager is only compatible with Airflow >= 3.0.0" -) -@pytest.mark.db_test -class TestAwsAuthManagerAuthenticationViews: - def test_login_redirect_to_identity_center(self, aws_app): - with aws_app.test_client() as client: - response = client.get("/login") - assert response.status_code == 302 - assert response.location.startswith("https://portal.sso.us-east-1.amazonaws.com/saml/assertion/") - - def test_logout_redirect_to_identity_center(self, aws_app): - with aws_app.test_client() as client: - response = client.post("/logout") - assert response.status_code == 302 - assert response.location.startswith("https://portal.sso.us-east-1.amazonaws.com/saml/logout/") - - def test_login_metadata_return_xml_file(self, aws_app): - with aws_app.test_client() as client: - response = client.get("/login_metadata") - assert response.status_code == 200 - assert response.headers["Content-Type"] == "text/xml" - - def test_login_callback_set_user_in_session(self): + mock_parse_remote.return_value = SAML_METADATA_PARSED + yield TestClient(create_app()) + + +class TestLoginRouter: + def test_login(self, test_client): + response = test_client.get("/auth/login", follow_redirects=False) + assert response.status_code == 307 + assert "location" in response.headers + assert response.headers["location"].startswith( + "https://portal.sso.us-east-1.amazonaws.com/saml/assertion/" + ) + + def test_login_callback_successful(self): with conf_vars( { ( @@ -104,18 +93,16 @@ def test_login_callback_set_user_in_session(self): } ): with ( + patch.object(OneLogin_Saml2_IdPMetadataParser, "parse_remote") as mock_parse_remote, patch( - "airflow.providers.amazon.aws.auth_manager.views.auth.OneLogin_Saml2_IdPMetadataParser" - ) as mock_parser, - patch( - "airflow.providers.amazon.aws.auth_manager.views.auth.AwsAuthManagerAuthenticationViews._init_saml_auth" + "airflow.providers.amazon.aws.auth_manager.router.login._init_saml_auth" ) as mock_init_saml_auth, patch( "airflow.providers.amazon.aws.auth_manager.avp.facade.AwsAuthManagerAmazonVerifiedPermissionsFacade.is_policy_store_schema_up_to_date" ) as mock_is_policy_store_schema_up_to_date, ): mock_is_policy_store_schema_up_to_date.return_value = True - mock_parser.parse_remote.return_value = SAML_METADATA_PARSED + mock_parse_remote.return_value = SAML_METADATA_PARSED auth = Mock() auth.is_authenticated.return_value = True @@ -126,16 +113,13 @@ def test_login_callback_set_user_in_session(self): "email": ["email"], } mock_init_saml_auth.return_value = auth - app = application.create_app(testing=True) - with app.test_client() as client: - response = client.get("/login_callback") - assert response.status_code == 302 - assert response.location == url_for("Airflow.index") - assert session["aws_user"] is not None - assert session["aws_user"].get_id() == "1" - assert session["aws_user"].get_name() == "user_id" - - def test_login_callback_raise_exception_if_errors(self): + client = TestClient(create_app()) + response = client.post("/auth/login_callback", follow_redirects=False) + assert response.status_code == 303 + assert "location" in response.headers + assert response.headers["location"].startswith("/webapp?token=") + + def test_login_callback_unsuccessful(self): with conf_vars( { ( @@ -146,28 +130,20 @@ def test_login_callback_raise_exception_if_errors(self): } ): with ( + patch.object(OneLogin_Saml2_IdPMetadataParser, "parse_remote") as mock_parse_remote, patch( - "airflow.providers.amazon.aws.auth_manager.views.auth.OneLogin_Saml2_IdPMetadataParser" - ) as mock_parser, - patch( - "airflow.providers.amazon.aws.auth_manager.views.auth.AwsAuthManagerAuthenticationViews._init_saml_auth" + "airflow.providers.amazon.aws.auth_manager.router.login._init_saml_auth" ) as mock_init_saml_auth, patch( "airflow.providers.amazon.aws.auth_manager.avp.facade.AwsAuthManagerAmazonVerifiedPermissionsFacade.is_policy_store_schema_up_to_date" ) as mock_is_policy_store_schema_up_to_date, ): mock_is_policy_store_schema_up_to_date.return_value = True - mock_parser.parse_remote.return_value = SAML_METADATA_PARSED + mock_parse_remote.return_value = SAML_METADATA_PARSED auth = Mock() auth.is_authenticated.return_value = False mock_init_saml_auth.return_value = auth - app = application.create_app(testing=True) - with app.test_client() as client: - with pytest.raises(AirflowException): - client.get("/login_callback") - - def test_logout_callback_raise_not_implemented_error(self, aws_app): - with aws_app.test_client() as client: - with pytest.raises(NotImplementedError): - client.get("/logout_callback") + client = TestClient(create_app()) + response = client.post("/auth/login_callback") + assert response.status_code == 500 diff --git a/providers/tests/amazon/aws/auth_manager/test_aws_auth_manager.py b/providers/tests/amazon/aws/auth_manager/test_aws_auth_manager.py index 797f96ee993b9..797d6e0c504f7 100644 --- a/providers/tests/amazon/aws/auth_manager/test_aws_auth_manager.py +++ b/providers/tests/amazon/aws/auth_manager/test_aws_auth_manager.py @@ -38,11 +38,7 @@ VariableDetails, ) from airflow.providers.amazon.aws.auth_manager.avp.entities import AvpEntities -from airflow.providers.amazon.aws.auth_manager.avp.facade import AwsAuthManagerAmazonVerifiedPermissionsFacade from airflow.providers.amazon.aws.auth_manager.aws_auth_manager import AwsAuthManager -from airflow.providers.amazon.aws.auth_manager.security_manager.aws_security_manager_override import ( - AwsSecurityManagerOverride, -) from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser from airflow.security.permissions import ( RESOURCE_AUDIT_LOG, @@ -54,7 +50,6 @@ from airflow.www.extensions.init_appbuilder import init_appbuilder from tests_common.test_utils.config import conf_vars -from tests_common.test_utils.www import check_content_in_response if TYPE_CHECKING: from airflow.auth.managers.base_auth_manager import ResourceMethod @@ -718,56 +713,9 @@ def test_filter_permitted_dag_ids(self, methods, user, auth_manager, test_user): auth_manager.avp_facade.get_batch_is_authorized_results.assert_called() assert result == {"dag_2"} - @patch("airflow.providers.amazon.aws.auth_manager.aws_auth_manager.url_for") - def test_get_url_login(self, mock_url_for, auth_manager): - auth_manager.get_url_login() - mock_url_for.assert_called_once_with("AwsAuthManagerAuthenticationViews.login") - - @patch("airflow.providers.amazon.aws.auth_manager.aws_auth_manager.url_for") - def test_get_url_logout(self, mock_url_for, auth_manager): - auth_manager.get_url_logout() - mock_url_for.assert_called_once_with("AwsAuthManagerAuthenticationViews.logout") - - @pytest.mark.db_test - def test_security_manager_return_default_security_manager(self, auth_manager_with_appbuilder): - assert isinstance(auth_manager_with_appbuilder.security_manager, AwsSecurityManagerOverride) + def test_get_url_login(self, auth_manager): + result = auth_manager.get_url_login() + assert result == "http://localhost:29091/auth/login" def test_get_cli_commands_return_cli_commands(self, auth_manager): assert len(auth_manager.get_cli_commands()) > 0 - - @pytest.mark.db_test - @patch( - "airflow.providers.amazon.aws.auth_manager.views.auth.conf.get_mandatory_value", return_value="test" - ) - def test_register_views(self, mock_get_mandatory_value, auth_manager_with_appbuilder): - from airflow.providers.amazon.aws.auth_manager.views.auth import AwsAuthManagerAuthenticationViews - - with patch.object(AwsAuthManagerAuthenticationViews, "idp_data"): - auth_manager_with_appbuilder.appbuilder.add_view_no_menu = Mock() - auth_manager_with_appbuilder.register_views() - auth_manager_with_appbuilder.appbuilder.add_view_no_menu.assert_called_once() - assert isinstance( - auth_manager_with_appbuilder.appbuilder.add_view_no_menu.call_args.args[0], - AwsAuthManagerAuthenticationViews, - ) - - @pytest.mark.db_test - @patch.object(AwsAuthManagerAmazonVerifiedPermissionsFacade, "get_batch_is_authorized_single_result") - @patch.object(AwsAuthManagerAmazonVerifiedPermissionsFacade, "get_batch_is_authorized_results") - @patch.object(AwsAuthManagerAmazonVerifiedPermissionsFacade, "is_authorized") - def test_aws_auth_manager_index( - self, - mock_is_authorized, - mock_get_batch_is_authorized_results, - mock_get_batch_is_authorized_single_result, - client_admin, - ): - """ - Load the index page using AWS auth manager. Mock all interactions with Amazon Verified Permissions. - """ - mock_is_authorized.return_value = True - mock_get_batch_is_authorized_results.return_value = [] - mock_get_batch_is_authorized_single_result.return_value = {"decision": "ALLOW"} - with client_admin.test_client() as client: - response = client.get("/login_callback", follow_redirects=True) - check_content_in_response("

DAGs

", response, 200)