fix: try002 for provider fab (#38801)

apache · Apr 7, 2024 · 6d3d207 · 6d3d207
1 parent 1533703
commit 6d3d207
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 8 deletions.
diff --git a/airflow/providers/fab/auth_manager/fab_auth_manager.py b/airflow/providers/fab/auth_manager/fab_auth_manager.py
@@ -44,7 +44,7 @@
     GroupCommand,
 )
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowConfigException, AirflowException
 from airflow.models import DagModel
 from airflow.providers.fab.auth_manager.cli_commands.definition import (
     ROLES_COMMANDS,
@@ -339,7 +339,7 @@ def security_manager(self) -> FabAirflowSecurityManagerOverride:
         sm_from_config = self.appbuilder.get_app.config.get("SECURITY_MANAGER_CLASS")
         if sm_from_config:
             if not issubclass(sm_from_config, FabAirflowSecurityManagerOverride):
-                raise Exception(
+                raise AirflowConfigException(
                     """Your CUSTOM_SECURITY_MANAGER must extend FabAirflowSecurityManagerOverride."""
                 )
             return sm_from_config(self.appbuilder)

diff --git a/airflow/providers/fab/auth_manager/security_manager/override.py b/airflow/providers/fab/auth_manager/security_manager/override.py
@@ -810,7 +810,7 @@ def _init_config(self):
         # LDAP Config
         if self.auth_type == AUTH_LDAP:
             if "AUTH_LDAP_SERVER" not in app.config:
-                raise Exception("No AUTH_LDAP_SERVER defined on config with AUTH_LDAP authentication type.")
+                raise ValueError("No AUTH_LDAP_SERVER defined on config with AUTH_LDAP authentication type.")
             app.config.setdefault("AUTH_LDAP_SEARCH", "")
             app.config.setdefault("AUTH_LDAP_SEARCH_FILTER", "")
             app.config.setdefault("AUTH_LDAP_APPEND_DOMAIN", "")

diff --git a/pyproject.toml b/pyproject.toml
@@ -386,9 +386,6 @@ combine-as-imports = true
 "airflow/providers/cncf/kubernetes/operators/pod.py" = ["TRY002"]
 # common.sql
 "airflow/providers/common/sql/hooks/sql.py" = ["TRY002"]
-# fab
-"airflow/providers/fab/auth_manager/fab_auth_manager.py" = ["TRY002"]
-"airflow/providers/fab/auth_manager/security_manager/override.py" = ["TRY002"]
 # google
 "airflow/providers/google/cloud/hooks/bigquery.py" = ["TRY002"]
 "airflow/providers/google/cloud/hooks/dataflow.py" = ["TRY002"]

diff --git a/tests/providers/fab/auth_manager/test_fab_auth_manager.py b/tests/providers/fab/auth_manager/test_fab_auth_manager.py
@@ -25,7 +25,7 @@
 from flask import Flask
 
 from airflow.auth.managers.models.resource_details import AccessView, DagAccessEntity, DagDetails
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowConfigException, AirflowException
 from airflow.providers.fab.auth_manager.fab_auth_manager import FabAuthManager
 from airflow.providers.fab.auth_manager.models import User
 from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride
@@ -424,7 +424,8 @@ class TestSecurityManager:
         flask_app.config["SECURITY_MANAGER_CLASS"] = TestSecurityManager
 
         with pytest.raises(
-            Exception, match="Your CUSTOM_SECURITY_MANAGER must extend FabAirflowSecurityManagerOverride."
+            AirflowConfigException,
+            match="Your CUSTOM_SECURITY_MANAGER must extend FabAirflowSecurityManagerOverride.",
         ):
             auth_manager_with_appbuilder.security_manager