refactor token creation logic to use methods in data models

Author: ecodina

providers/keycloak/src/airflow/providers/keycloak/auth_manager/datamodels/token.py

@@ -22,6 +22,10 @@
 from pydantic import Field, RootModel, model_validator
 
 from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel
+from airflow.providers.keycloak.auth_manager.services.token import (
+    create_client_credentials_token,
+    create_token_for,
+)
 
 
 class TokenResponse(BaseModel):
@@ -37,6 +41,12 @@ class TokenPasswordBody(StrictBaseModel):
     username: str = Field()
     password: str = Field()
 
+    def create_token(self, expiration_time_in_seconds: int) -> str:
+        """Create token using password grant."""
+        return create_token_for(
+            self.username, self.password, expiration_time_in_seconds=expiration_time_in_seconds
+        )
+
 
 class TokenClientCredentialsBody(StrictBaseModel):
     """Client Credentials Grant Token serializer for post bodies."""
@@ -45,6 +55,12 @@ class TokenClientCredentialsBody(StrictBaseModel):
     client_id: str = Field()
     client_secret: str = Field()
 
+    def create_token(self, expiration_time_in_seconds: int) -> str:
+        """Create token using client credentials grant."""
+        return create_client_credentials_token(
+            self.client_id, self.client_secret, expiration_time_in_seconds=expiration_time_in_seconds
+        )
+
 
 TokenUnion = Annotated[
     TokenPasswordBody | TokenClientCredentialsBody,

providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/token.py

@@ -26,14 +26,9 @@
 from airflow.configuration import conf
 from airflow.providers.keycloak.auth_manager.datamodels.token import (
     TokenBody,
-    TokenClientCredentialsBody,
     TokenPasswordBody,
     TokenResponse,
 )
-from airflow.providers.keycloak.auth_manager.services.token import (
-    create_client_credentials_token,
-    create_token_for,
-)
 
 log = logging.getLogger(__name__)
 token_router = AirflowRouter(tags=["KeycloakAuthManagerToken"])
@@ -44,16 +39,10 @@
     status_code=status.HTTP_201_CREATED,
     responses=create_openapi_http_exception_doc([status.HTTP_400_BAD_REQUEST, status.HTTP_401_UNAUTHORIZED]),
 )
-def create_token(
-    body: TokenBody,
-) -> TokenResponse:
-    credentials = body.root
-    if isinstance(credentials, TokenPasswordBody):
-        token = create_token_for(credentials.username, credentials.password)
-    elif isinstance(credentials, TokenClientCredentialsBody):
-        token = create_client_credentials_token(credentials.client_id, credentials.client_secret)
-    else:
-        raise ValueError("Unsupported grant_type")
+def create_token(body: TokenBody) -> TokenResponse:
+    token = body.root.create_token(
+        expiration_time_in_seconds=int(conf.getint("api_auth", "jwt_expiration_time"))
+    )
     return TokenResponse(access_token=token)
 
 
@@ -63,9 +52,7 @@ def create_token(
     responses=create_openapi_http_exception_doc([status.HTTP_400_BAD_REQUEST, status.HTTP_401_UNAUTHORIZED]),
 )
 def create_token_cli(body: TokenPasswordBody) -> TokenResponse:
-    token = create_token_for(
-        body.username,
-        body.password,
-        expiration_time_in_seconds=int(conf.getint("api_auth", "jwt_cli_expiration_time")),
+    token = body.create_token(
+        expiration_time_in_seconds=int(conf.getint("api_auth", "jwt_cli_expiration_time"))
     )
     return TokenResponse(access_token=token)

providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_token.py

@@ -41,7 +41,7 @@ class TestTokenRouter:
             ("api_auth", "jwt_expiration_time"): "10",
         }
     )
-    @patch("airflow.providers.keycloak.auth_manager.routes.token.create_token_for")
+    @patch("airflow.providers.keycloak.auth_manager.datamodels.token.create_token_for")
     def test_create_token_password_grant(self, mock_create_token_for, client, body):
         mock_create_token_for.return_value = self.token
         response = client.post(
@@ -58,7 +58,7 @@ def test_create_token_password_grant(self, mock_create_token_for, client, body):
             ("api_auth", "jwt_expiration_time"): "10",
         }
     )
-    @patch("airflow.providers.keycloak.auth_manager.routes.token.create_token_for")
+    @patch("airflow.providers.keycloak.auth_manager.datamodels.token.create_token_for")
     def test_create_token_cli(self, mock_create_token_for, client):
         mock_create_token_for.return_value = self.token
         response = client.post(
@@ -74,7 +74,7 @@ def test_create_token_cli(self, mock_create_token_for, client):
             ("api_auth", "jwt_expiration_time"): "10",
         }
     )
-    @patch("airflow.providers.keycloak.auth_manager.routes.token.create_client_credentials_token")
+    @patch("airflow.providers.keycloak.auth_manager.datamodels.token.create_client_credentials_token")
     def test_create_token_client_credentials(self, mock_create_client_credentials_token, client):
         mock_create_client_credentials_token.return_value = self.token
         response = client.post(
@@ -88,7 +88,9 @@ def test_create_token_client_credentials(self, mock_create_client_credentials_to
 
         assert response.status_code == 201
         assert response.json() == {"access_token": self.token}
-        mock_create_client_credentials_token.assert_called_once_with("client_id", "client_secret")
+        mock_create_client_credentials_token.assert_called_once_with(
+            "client_id", "client_secret", expiration_time_in_seconds=10
+        )
 
     @pytest.mark.parametrize(
         "body",
@@ -107,7 +109,7 @@ def test_create_token_client_credentials(self, mock_create_client_credentials_to
             ("api_auth", "jwt_expiration_time"): "10",
         }
     )
-    @patch("airflow.providers.keycloak.auth_manager.routes.token.create_client_credentials_token")
+    @patch("airflow.providers.keycloak.auth_manager.datamodels.token.create_client_credentials_token")
     def test_create_token_invalid_body(self, mock_create_client_credentials_token, client, body):
         mock_create_client_credentials_token.return_value = self.token
         response = client.post(