fix tests

Author: loveyana

tests/a2a/__init__.py

@@ -89,7 +89,9 @@ def test_initialization(self, credential_service):
         assert credential_service._credentials == {}
 
     @pytest.mark.asyncio
-    async def test_set_and_get_credential(self, credential_service, sample_auth_credential):
+    async def test_set_and_get_credential(
+        self, credential_service, sample_auth_credential
+    ):
         """Test setting and getting credentials."""
         # Set credential
         await credential_service.set_credential(
@@ -234,7 +236,11 @@ async def test_save_credential_via_adk_interface(
 
     @pytest.mark.asyncio
     async def test_load_credential_via_adk_interface(
-        self, credential_service, mock_callback_context, mock_auth_config, sample_auth_credential
+        self,
+        credential_service,
+        mock_callback_context,
+        mock_auth_config,
+        sample_auth_credential,
     ):
         """Test loading credential via ADK BaseCredentialService interface."""
         # First set a credential
@@ -356,4 +362,3 @@ async def test_credential_isolation_between_apps(self, credential_service):
 
         assert retrieved1.http.credentials.token == "app1_token"
         assert retrieved2.http.credentials.token == "app2_token"
-

tests/auth/test_credential_service.py

@@ -15,11 +15,10 @@
 """Unit tests for A2A authentication middleware."""
 
 import pytest
-from unittest.mock import Mock, AsyncMock, patch
+from unittest.mock import Mock, patch
 from starlette.applications import Starlette
 from starlette.requests import Request
 from starlette.responses import Response
-from starlette.testclient import TestClient
 
 from veadk.a2a.ve_middlewares import A2AAuthMiddleware, build_a2a_auth_middleware
 from veadk.auth.credential_service import VeCredentialService
@@ -75,17 +74,19 @@ def test_middleware_initialization(self, credential_service, mock_identity_clien
     def test_middleware_default_identity_client(self, credential_service):
         """Test middleware uses global identity client when not provided."""
         app = Starlette()
-        
-        with patch("veadk.a2a.ve_a2a_middlewares.get_default_identity_client") as mock_get_client:
+
+        with patch(
+            "veadk.a2a.ve_middlewares.get_default_identity_client"
+        ) as mock_get_client:
             mock_client = Mock()
             mock_get_client.return_value = mock_client
-            
+
             middleware = A2AAuthMiddleware(
                 app=app,
                 app_name="test_app",
                 credential_service=credential_service,
             )
-            
+
             mock_get_client.assert_called_once()
             assert middleware.identity_client == mock_client
 
@@ -104,7 +105,7 @@ def test_extract_token_from_header_with_bearer(self, credential_service):
         mock_request.headers = {"Authorization": "Bearer test_token_123"}
 
         token, has_prefix = middleware._extract_token(mock_request)
-        
+
         assert token == "test_token_123"
         assert has_prefix is True
 
@@ -122,7 +123,7 @@ def test_extract_token_from_header_without_bearer(self, credential_service):
         mock_request.headers = {"Authorization": "test_token_123"}
 
         token, has_prefix = middleware._extract_token(mock_request)
-        
+
         assert token == "test_token_123"
         assert has_prefix is False
 
@@ -174,7 +175,9 @@ async def test_dispatch_with_valid_jwt_token(
         mock_workload_token = Mock()
         mock_workload_token.workload_access_token = "workload_token_123"
         mock_workload_token.expires_at = 1234567890
-        mock_identity_client.get_workload_access_token.return_value = mock_workload_token
+        mock_identity_client.get_workload_access_token.return_value = (
+            mock_workload_token
+        )
 
         middleware = A2AAuthMiddleware(
             app=app,
@@ -196,10 +199,14 @@ async def mock_call_next(request):
             return Response("OK", status_code=200)
 
         # Execute dispatch
-        with patch("veadk.a2a.ve_a2a_middlewares.extract_delegation_chain_from_jwt") as mock_extract:
+        with patch(
+            "veadk.a2a.ve_middlewares.extract_delegation_chain_from_jwt"
+        ) as mock_extract:
             mock_extract.return_value = ("user123", ["agent1"])
 
-            with patch("veadk.a2a.ve_a2a_middlewares.build_auth_config") as mock_build_config:
+            with patch(
+                "veadk.a2a.ve_middlewares.build_auth_config"
+            ) as mock_build_config:
                 mock_auth_config = Mock()
                 mock_auth_config.exchanged_auth_credential = Mock()
                 mock_build_config.return_value = mock_auth_config
@@ -228,7 +235,9 @@ async def test_dispatch_with_tip_token(
         mock_workload_token = Mock()
         mock_workload_token.workload_access_token = "workload_token_from_tip"
         mock_workload_token.expires_at = 1234567890
-        mock_identity_client.get_workload_access_token.return_value = mock_workload_token
+        mock_identity_client.get_workload_access_token.return_value = (
+            mock_workload_token
+        )
 
         middleware = A2AAuthMiddleware(
             app=app,
@@ -252,10 +261,14 @@ async def mock_call_next(request):
             return Response("OK", status_code=200)
 
         # Execute dispatch
-        with patch("veadk.a2a.ve_a2a_middlewares.extract_delegation_chain_from_jwt") as mock_extract:
+        with patch(
+            "veadk.a2a.ve_middlewares.extract_delegation_chain_from_jwt"
+        ) as mock_extract:
             mock_extract.return_value = ("user123", ["agent1"])
 
-            with patch("veadk.a2a.ve_a2a_middlewares.build_auth_config") as mock_build_config:
+            with patch(
+                "veadk.a2a.ve_middlewares.build_auth_config"
+            ) as mock_build_config:
                 mock_auth_config = Mock()
                 mock_auth_config.exchanged_auth_credential = Mock()
                 mock_build_config.return_value = mock_auth_config
@@ -264,8 +277,7 @@ async def mock_call_next(request):
 
         # Verify TIP token was used for workload token exchange
         mock_identity_client.get_workload_access_token.assert_called_once_with(
-            user_token=tip_token,
-            user_id="user123"
+            user_token=tip_token, user_id="user123"
         )
 
         # Verify workload token was set
@@ -285,7 +297,9 @@ def test_build_middleware_basic(self, credential_service):
         assert middleware_class is not None
         assert issubclass(middleware_class, A2AAuthMiddleware)
 
-    def test_build_middleware_with_all_params(self, credential_service, mock_identity_client):
+    def test_build_middleware_with_all_params(
+        self, credential_service, mock_identity_client
+    ):
         """Test building middleware with all parameters."""
         middleware_class = build_a2a_auth_middleware(
             app_name="test_app",
@@ -305,4 +319,3 @@ def test_build_middleware_with_all_params(self, credential_service, mock_identit
         assert instance.token_param == "access_token"
         assert instance.credential_key == "custom_key"
         assert instance.identity_client == mock_identity_client
-

tests/a2a/test_ve_a2a_middlewares.py tests/test_ve_a2a_middlewares.pytests/a2a/test_ve_a2a_middlewares.py renamed to tests/test_ve_a2a_middlewares.py

@@ -58,7 +58,7 @@ class A2AAuthMiddleware(BaseHTTPMiddleware):
     Examples:
         ```python
         from fastapi import FastAPI
-        from veadk.a2a.ve_a2a_middlewares import build_a2a_auth_middleware
+        from veadk.a2a.ve_middlewares import build_a2a_auth_middleware
         from veadk.auth.credential_service import VeCredentialService
 
         app = FastAPI()
@@ -273,7 +273,7 @@ def build_a2a_auth_middleware(
     Examples:
         ```python
         from fastapi import FastAPI, Request
-        from veadk.a2a.ve_a2a_middlewares import build_a2a_auth_middleware
+        from veadk.a2a.ve_middlewares import build_a2a_auth_middleware
         from veadk.auth.credential_service import VeCredentialService
         from veadk.integrations.ve_identity import IdentityClient