Fix a typo and update tests with a rename

dmandar · dmandar · commit d94df5a23368 · 2025-06-16T16:35:24.000-07:00
diff --git a/src/a2a/client/auth/credentials.py b/src/a2a/client/auth/credentials.py
@@ -25,7 +25,6 @@ class InMemoryContextCredentialStore(CredentialService):
     """
 
     def __init__(self) -> None:
-        # {session_id: {scheme_name: credential}}
         self._store: dict[str, dict[str, str]] = {}
 
     async def get_credentials(
@@ -38,7 +37,7 @@ async def get_credentials(
         session_id = context.state['sessionId']
         return self._store.get(session_id, {}).get(security_scheme_name)
 
-    async def set_credential(
+    async def set_credentials(
         self, session_id: str, security_scheme_name: str, credential: str
     ) -> None:
         """Method to populate the store."""
diff --git a/src/a2a/client/middleware.py b/src/a2a/client/middleware.py
@@ -1,10 +1,11 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from collections.abc import MutableMapping
 from typing import TYPE_CHECKING, Any
 
 from pydantic import BaseModel, Field
 
-
 if TYPE_CHECKING:
     from a2a.types import AgentCard
 
@@ -31,8 +32,8 @@ async def intercept(
         method_name: str,
         request_payload: dict[str, Any],
         http_kwargs: dict[str, Any],
-        agent_card: 'AgentCard | None',
-        context: 'ClientCallContext | None',
+        agent_card: AgentCard | None,
+        context: ClientCallContext | None,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
         """
         Intercepts a client call before the request is sent.
diff --git a/tests/client/test_auth_middleware.py b/tests/client/test_auth_middleware.py
@@ -6,18 +6,10 @@
 
 from a2a.client import A2AClient, ClientCallContext, ClientCallInterceptor
 from a2a.client.auth import AuthInterceptor, InMemoryContextCredentialStore
-from a2a.types import (
-    APIKeySecurityScheme,
-    AgentCapabilities,
-    AgentCard,
-    AuthorizationCodeOAuthFlow,
-    In,
-    OAuth2SecurityScheme,
-    OAuthFlows,
-    OpenIdConnectSecurityScheme,
-    SecurityScheme,
-    SendMessageRequest,
-)
+from a2a.types import (AgentCapabilities, AgentCard, APIKeySecurityScheme,
+                       AuthorizationCodeOAuthFlow, In, OAuth2SecurityScheme,
+                       OAuthFlows, OpenIdConnectSecurityScheme, SecurityScheme,
+                       SendMessageRequest)
 
 
 # A simple mock interceptor for testing basic middleware functionality
@@ -66,7 +58,7 @@ async def test_client_with_simple_interceptor():
                 'messageId': 'response-msg',
                 'role': 'agent',
                 'parts': [],
-            },  # Example result
+            },
         }
         respx.post(test_url).mock(
             return_value=httpx.Response(200, json=minimal_success_response)
@@ -97,24 +89,24 @@ async def test_client_with_simple_interceptor():
 async def test_in_memory_context_credential_store():
     """
     Tests the functionality of the InMemoryContextCredentialStore to ensure
-    it correctly stores and retrieves credentials based on contextId.
+    it correctly stores and retrieves credentials based on sessionId.
     """
     # Arrange
     store = InMemoryContextCredentialStore()
-    context_id = 'test-context-123'
+    session_id = 'test-session-123'
     scheme_name = 'test-scheme'
     credential = 'test-token'
 
     # Act
-    await store.set_credentials(context_id, scheme_name, credential)
+    await store.set_credentials(session_id, scheme_name, credential)
 
     # Assert: Successful retrieval
-    context = ClientCallContext(state={'contextId': context_id})
+    context = ClientCallContext(state={'sessionId': session_id})
     retrieved_credential = await store.get_credentials(scheme_name, context)
     assert retrieved_credential == credential
 
-    # Assert: Retrieval with wrong context returns None
-    wrong_context = ClientCallContext(state={'contextId': 'wrong-context'})
+    # Assert: Retrieval with wrong session ID returns None
+    wrong_context = ClientCallContext(state={'sessionId': 'wrong-session'})
     retrieved_credential_wrong = await store.get_credentials(
         scheme_name, wrong_context
     )
@@ -124,7 +116,7 @@ async def test_in_memory_context_credential_store():
     retrieved_credential_none = await store.get_credentials(scheme_name, None)
     assert retrieved_credential_none is None
 
-    # Assert: Retrieval with context but no contextId returns None
+    # Assert: Retrieval with context but no sessionId returns None
     empty_context = ClientCallContext(state={})
     retrieved_credential_empty = await store.get_credentials(
         scheme_name, empty_context
@@ -140,15 +132,21 @@ async def test_auth_interceptor_with_api_key():
     """
     # Arrange
     test_url = 'http://apikey-agent.com/rpc'
-    context_id = 'user-session-2'
+    session_id = 'user-session-2'
     scheme_name = 'apiKeyAuth'
     api_key = 'secret-api-key'
 
     cred_store = InMemoryContextCredentialStore()
-    await cred_store.set_credentials(context_id, scheme_name, api_key)
+    await cred_store.set_credentials(session_id, scheme_name, api_key)
 
     auth_interceptor = AuthInterceptor(credential_service=cred_store)
 
+    api_key_scheme_params = {
+        'type': 'apiKey',
+        'name': 'X-API-Key',
+        'in': In.header,
+    }
+
     agent_card = AgentCard(
         url=test_url,
         name='ApiKeyBot',
@@ -161,82 +159,11 @@ async def test_auth_interceptor_with_api_key():
         security=[{scheme_name: []}],
         securitySchemes={
             scheme_name: SecurityScheme(
-                root=APIKeySecurityScheme(
-                    name='X-API-Key', in_=In.header, type='apiKey'
-                )
+                root=APIKeySecurityScheme(**api_key_scheme_params)
             )
         },
     )
 
-    async with httpx.AsyncClient() as http_client:
-        client = A2AClient(
-            httpx_client=http_client,
-            agent_card=agent_card,
-            interceptors=[auth_interceptor],
-        )
-
-        respx.post(test_url).mock(return_value=httpx.Response(200, json={}))
-
-        # Act
-        context = ClientCallContext(state={'contextId': context_id})
-        await client.send_message(
-            request=SendMessageRequest(
-                id='1',
-                params={
-                    'message': {
-                        'messageId': 'msg1',
-                        'role': 'user',
-                        'parts': [],
-                    }
-                },
-            ),
-            context=context,
-        )
-
-        # Assert
-        assert len(respx.calls) == 1
-        request = respx.calls.last.request
-        assert 'x-api-key' in request.headers
-        assert request.headers['x-api-key'] == api_key
-
-
-@pytest.mark.asyncio
-@respx.mock
-async def test_auth_interceptor_with_api_key():
-    """
-    Tests the authentication flow with an API key in the header.
-    """
-    # Arrange
-    test_url = 'http://apikey-agent.com/rpc'
-    context_id = 'user-session-2'
-    scheme_name = 'apiKeyAuth'
-    api_key = 'secret-api-key'
-
-    cred_store = InMemoryContextCredentialStore()
-    await cred_store.set_credentials(context_id, scheme_name, api_key)
-
-    auth_interceptor = AuthInterceptor(credential_service=cred_store)
-
-    # Use the alias 'in' for instantiation, as confirmed by debug output
-    api_key_scheme_instance = APIKeySecurityScheme(
-        name='X-API-Key', **{'in': In.header}
-    )
-
-    agent_card = AgentCard(
-        url=test_url,
-        name='ApiKeyBot',
-        description='A bot that requires an API Key',
-        version='1.0',
-        defaultInputModes=[],
-        defaultOutputModes=[],
-        skills=[],
-        capabilities=AgentCapabilities(),
-        security=[{scheme_name: []}],
-        securitySchemes={
-            scheme_name: SecurityScheme(root=api_key_scheme_instance)
-        },
-    )
-
     async with httpx.AsyncClient() as http_client:
         client = A2AClient(
             httpx_client=http_client,
@@ -259,7 +186,7 @@ async def test_auth_interceptor_with_api_key():
         )
 
         # Act
-        context = ClientCallContext(state={'contextId': context_id})
+        context = ClientCallContext(state={'sessionId': session_id})
         await client.send_message(
             request=SendMessageRequest(
                 id='1',
@@ -289,16 +216,15 @@ async def test_auth_interceptor_with_oauth2_scheme():
     Ensures it correctly sets the Authorization: Bearer <token> header.
     """
     test_url = 'http://oauth-agent.com/rpc'
-    context_id = 'user-session-oauth'
+    session_id = 'user-session-oauth'
     scheme_name = 'myOAuthScheme'
     access_token = 'secret-oauth-access-token'
 
     cred_store = InMemoryContextCredentialStore()
-    await cred_store.set_credentials(context_id, scheme_name, access_token)
+    await cred_store.set_credentials(session_id, scheme_name, access_token)
 
     auth_interceptor = AuthInterceptor(credential_service=cred_store)
 
-    # Define a minimal OAuth2 flow
     oauth_flows = OAuthFlows(
         authorizationCode=AuthorizationCodeOAuthFlow(
             authorizationUrl='http://provider.com/auth',
@@ -346,7 +272,7 @@ async def test_auth_interceptor_with_oauth2_scheme():
         )
 
         # Act
-        context = ClientCallContext(state={'contextId': context_id})
+        context = ClientCallContext(state={'sessionId': session_id})
         await client.send_message(
             request=SendMessageRequest(
                 id='oauth_test_1',
@@ -377,12 +303,12 @@ async def test_auth_interceptor_with_oidc_scheme():
     """
     # Arrange
     test_url = 'http://oidc-agent.com/rpc'
-    context_id = 'user-session-oidc'
+    session_id = 'user-session-oidc'
     scheme_name = 'myOidcScheme'
-    id_token = 'secret-oidc-id-token'  # Or access_token
+    id_token = 'secret-oidc-id-token'
 
     cred_store = InMemoryContextCredentialStore()
-    await cred_store.set_credentials(context_id, scheme_name, id_token)
+    await cred_store.set_credentials(session_id, scheme_name, id_token)
 
     auth_interceptor = AuthInterceptor(credential_service=cred_store)
 
@@ -397,7 +323,7 @@ async def test_auth_interceptor_with_oidc_scheme():
         capabilities=AgentCapabilities(),
         security=[
             {scheme_name: []}
-        ],  # Security requirement referencing the scheme
+        ],
         securitySchemes={
             scheme_name: SecurityScheme(
                 root=OpenIdConnectSecurityScheme(
@@ -430,7 +356,7 @@ async def test_auth_interceptor_with_oidc_scheme():
         )
 
         # Act
-        context = ClientCallContext(state={'contextId': context_id})
+        context = ClientCallContext(state={'sessionId': session_id})
         await client.send_message(
             request=SendMessageRequest(
                 id='oidc_test_1',
@@ -449,4 +375,4 @@ async def test_auth_interceptor_with_oidc_scheme():
         assert len(respx.calls) == 1
         request_sent = respx.calls.last.request
         assert 'Authorization' in request_sent.headers
-        assert request_sent.headers['Authorization'] == f'Bearer {id_token}'
+        assert request_sent.headers['Authorization'] == f'Bearer {id_token}'
