Add tests related to server side extensions.

Author: mikeas1

tests/extensions/test_common.py

@@ -0,0 +1,38 @@
+from a2a.extensions.common import find_extension_by_uri
+from a2a.types import AgentCard, AgentExtension, AgentCapabilities
+
+
+def test_find_extension_by_uri():
+    ext1 = AgentExtension(uri='foo', name='Foo', description='The Foo extension')
+    ext2 = AgentExtension(uri='bar', name='Bar', description='The Bar extension')
+    card = AgentCard(
+        agent_id='test-agent',
+        name='Test Agent',
+        description='Test Agent Description',
+        version='1.0',
+        url='http://test.com',
+        skills=[],
+        defaultInputModes=['text/plain'],
+        defaultOutputModes=['text/plain'],
+        capabilities=AgentCapabilities(extensions=[ext1, ext2]),
+    )
+
+    assert find_extension_by_uri(card, 'foo') == ext1
+    assert find_extension_by_uri(card, 'bar') == ext2
+    assert find_extension_by_uri(card, 'baz') is None
+
+
+def test_find_extension_by_uri_no_extensions():
+    card = AgentCard(
+        agent_id='test-agent',
+        name='Test Agent',
+        description='Test Agent Description',
+        version='1.0',
+        url='http://test.com',
+        skills=[],
+        defaultInputModes=['text/plain'],
+        defaultOutputModes=['text/plain'],
+        capabilities=AgentCapabilities(extensions=None),
+    )
+
+    assert find_extension_by_uri(card, 'foo') is None

tests/server/agent_execution/test_context.py

@@ -5,6 +5,7 @@
 import pytest
 
 from a2a.server.agent_execution import RequestContext
+from a2a.server.context import ServerCallContext
 from a2a.types import (
     Message,
     MessageSendParams,
@@ -262,3 +263,16 @@ def test_init_with_context_id_and_existing_context_id_match(
 
         assert context.context_id == mock_task.contextId
         assert context.current_task == mock_task
+
+    def test_extension_handling(self):
+        """Test extension handling in RequestContext."""
+        call_context = ServerCallContext(requested_extensions={'foo', 'bar'})
+        context = RequestContext(call_context=call_context)
+
+        assert context.requested_extensions == {'foo', 'bar'}
+
+        context.add_activated_extension('foo')
+        assert call_context.activated_extensions == {'foo'}
+
+        context.add_activated_extension('baz')
+        assert call_context.activated_extensions == {'foo', 'baz'}

tests/server/apps/jsonrpc/test_jsonrpc_app.py

@@ -1,23 +1,36 @@
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
-
+from starlette.applications import Starlette
+from starlette.testclient import TestClient
 
 # Attempt to import StarletteBaseUser, fallback to MagicMock if not available
 try:
     from starlette.authentication import BaseUser as StarletteBaseUser
 except ImportError:
     StarletteBaseUser = MagicMock()  # type: ignore
 
+from a2a.extensions.common import HTTP_EXTENSION_HEADER
 from a2a.server.apps.jsonrpc.jsonrpc_app import (
-    JSONRPCApplication,  # Still needed for JSONRPCApplication default constructor arg
+    JSONRPCApplication,
     StarletteUserProxy,
 )
+from a2a.server.apps.jsonrpc.starlette_app import A2AStarletteApplication
+from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.request_handler import (
-    RequestHandler,  # For mock spec
+    RequestHandler,
+)  # For mock spec
+from a2a.types import (
+    AgentCapabilities,
+    AgentCard,
+    Message,
+    MessageSendParams,
+    Role,
+    SendMessageRequest,
+    SendMessageResponse,
+    SendMessageSuccessResponse,
+    TextPart,
 )
-from a2a.types import AgentCard  # For mock spec
-
 
 # --- StarletteUserProxy Tests ---
 
@@ -69,6 +82,7 @@ def test_jsonrpc_app_build_method_abstract_raises_typeerror(
         mock_agent_card.url = 'http://mockurl.com'
         # Ensure 'supportsAuthenticatedExtendedCard' attribute exists
         mock_agent_card.supportsAuthenticatedExtendedCard = False
+        mock_agent_card.capabilities = AgentCapabilities(streaming=True)
 
         # This will fail at definition time if an abstract method is not implemented
         with pytest.raises(
@@ -86,5 +100,149 @@ def some_other_method(self):
             )
 
 
+class TestJSONRPCExtensions:
+    @pytest.fixture
+    def mock_handler(self):
+        handler = AsyncMock(spec=RequestHandler)
+        handler.on_message_send.return_value = SendMessageResponse(
+            root=SendMessageSuccessResponse(
+                id='1',
+                result=Message(
+                    messageId='test',
+                    role=Role.agent,
+                    parts=[TextPart(text='response message')],
+                ),
+            )
+        )
+        return handler
+
+    @pytest.fixture
+    def test_app(self, mock_handler):
+        mock_agent_card = MagicMock(spec=AgentCard)
+        mock_agent_card.url = 'http://mockurl.com'
+        mock_agent_card.supportsAuthenticatedExtendedCard = False
+
+        return A2AStarletteApplication(
+            agent_card=mock_agent_card, http_handler=mock_handler
+        )
+
+    @pytest.fixture
+    def client(self, test_app):
+        return TestClient(test_app.build())
+
+    def test_request_with_single_extension(self, client, mock_handler):
+        headers = {HTTP_EXTENSION_HEADER: 'foo'}
+        response = client.post(
+            '/',
+            headers=headers,
+            json=SendMessageRequest(
+                id='1',
+                params=MessageSendParams(
+                    message=Message(
+                        messageId='1',
+                        role=Role.user,
+                        parts=[TextPart(text='hi')],
+                    )
+                ),
+            ).model_dump(),
+        )
+        response.raise_for_status()
+
+        mock_handler.on_message_send.assert_called_once()
+        call_context = mock_handler.on_message_send.call_args[0][1]
+        assert isinstance(call_context, ServerCallContext)
+        assert call_context.requested_extensions == {'foo'}
+
+    def test_request_with_comma_separated_extensions(
+        self, client, mock_handler
+    ):
+        headers = {HTTP_EXTENSION_HEADER: 'foo, bar'}
+        response = client.post(
+            '/',
+            headers=headers,
+            json=SendMessageRequest(
+                id='1',
+                params=MessageSendParams(
+                    message=Message(
+                        messageId='1',
+                        role=Role.user,
+                        parts=[TextPart(text='hi')],
+                    )
+                ),
+            ).model_dump(),
+        )
+        response.raise_for_status()
+
+        mock_handler.on_message_send.assert_called_once()
+        call_context = mock_handler.on_message_send.call_args[0][1]
+        assert call_context.requested_extensions == {'foo', 'bar'}
+
+    def test_request_with_multiple_extension_headers(
+        self, client, mock_handler
+    ):
+        headers = [
+            (HTTP_EXTENSION_HEADER, 'foo'),
+            (HTTP_EXTENSION_HEADER, 'bar'),
+        ]
+        response = client.post(
+            '/',
+            headers=headers,
+            json=SendMessageRequest(
+                id='1',
+                params=MessageSendParams(
+                    message=Message(
+                        messageId='1',
+                        role=Role.user,
+                        parts=[TextPart(text='hi')],
+                    )
+                ),
+            ).model_dump(),
+        )
+        response.raise_for_status()
+
+        mock_handler.on_message_send.assert_called_once()
+        call_context = mock_handler.on_message_send.call_args[0][1]
+        assert call_context.requested_extensions == {'foo', 'bar'}
+
+    def test_response_with_activated_extensions(self, client, mock_handler):
+        def side_effect(request, context: ServerCallContext):
+            context.activated_extensions.add('foo')
+            context.activated_extensions.add('baz')
+            return SendMessageResponse(
+                root=SendMessageSuccessResponse(
+                    id='1',
+                    result=Message(
+                        messageId='test',
+                        role=Role.agent,
+                        parts=[TextPart(text='response message')],
+                    ),
+                )
+            )
+
+        mock_handler.on_message_send.side_effect = side_effect
+
+        response = client.post(
+            '/',
+            json=SendMessageRequest(
+                id='1',
+                params=MessageSendParams(
+                    message=Message(
+                        messageId='1',
+                        role=Role.user,
+                        parts=[TextPart(text='hi')],
+                    )
+                ),
+            ).model_dump(),
+        )
+        response.raise_for_status()
+
+        assert response.status_code == 200
+        assert HTTP_EXTENSION_HEADER in response.headers
+        assert set(response.headers[HTTP_EXTENSION_HEADER].split(', ')) == {
+            'foo',
+            'baz',
+        }
+
+
 if __name__ == '__main__':
     pytest.main([__file__])

tests/server/request_handlers/test_grpc_handler.py

@@ -1,14 +1,15 @@
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import grpc
 import pytest
 
 from a2a import types
+from a2a.extensions.common import HTTP_EXTENSION_HEADER
 from a2a.grpc import a2a_pb2
+from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers import GrpcHandler, RequestHandler
 from a2a.utils.errors import ServerError
 
-
 # --- Fixtures ---
 
 
@@ -21,6 +22,8 @@ def mock_request_handler() -> AsyncMock:
 def mock_grpc_context() -> AsyncMock:
     context = AsyncMock(spec=grpc.aio.ServicerContext)
     context.abort = AsyncMock()
+    context.invocation_metadata = MagicMock()
+    context.set_trailing_metadata = MagicMock()
     return context
 
 
@@ -279,3 +282,88 @@ async def test_abort_context_error_mapping(
     call_args, _ = mock_grpc_context.abort.call_args
     assert call_args[0] == grpc_status_code
     assert error_message_part in call_args[1]
+
+
+@pytest.mark.asyncio
+class TestGrpcExtensions:
+    @patch(
+        'a2a.server.request_handlers.grpc_handler.DefaultCallContextBuilder.build'
+    )
+    async def test_send_message_with_extensions(
+        self,
+        mock_build,
+        grpc_handler: GrpcHandler,
+        mock_request_handler: AsyncMock,
+        mock_grpc_context: AsyncMock,
+    ):
+        mock_build.return_value = ServerCallContext(
+            requested_extensions={'foo', 'bar'}
+        )
+
+        def side_effect(request, context: ServerCallContext):
+            context.activated_extensions.add('foo')
+            context.activated_extensions.add('baz')
+            return types.Task(
+                id='task-1',
+                contextId='ctx-1',
+                status=types.TaskStatus(state=types.TaskState.completed),
+            )
+
+        mock_request_handler.on_message_send.side_effect = side_effect
+
+        await grpc_handler.SendMessage(
+            a2a_pb2.SendMessageRequest(), mock_grpc_context
+        )
+
+        mock_request_handler.on_message_send.assert_awaited_once()
+        call_context = mock_request_handler.on_message_send.call_args[0][1]
+        assert isinstance(call_context, ServerCallContext)
+        assert call_context.requested_extensions == {'foo', 'bar'}
+
+        mock_grpc_context.set_trailing_metadata.assert_called_once_with(
+            [(HTTP_EXTENSION_HEADER, 'foo'), (HTTP_EXTENSION_HEADER, 'baz')]
+        )
+
+    @patch(
+        'a2a.server.request_handlers.grpc_handler.DefaultCallContextBuilder.build'
+    )
+    async def test_send_streaming_message_with_extensions(
+        self,
+        mock_build,
+        grpc_handler: GrpcHandler,
+        mock_request_handler: AsyncMock,
+        mock_grpc_context: AsyncMock,
+    ):
+        mock_build.return_value = ServerCallContext(
+            requested_extensions={'foo', 'bar'}
+        )
+
+        async def side_effect(request, context: ServerCallContext):
+            context.activated_extensions.add('foo')
+            context.activated_extensions.add('baz')
+            yield types.Task(
+                id='task-1',
+                contextId='ctx-1',
+                status=types.TaskStatus(state=types.TaskState.working),
+            )
+
+        mock_request_handler.on_message_send_stream.side_effect = side_effect
+
+        results = [
+            result
+            async for result in grpc_handler.SendStreamingMessage(
+                a2a_pb2.SendMessageRequest(), mock_grpc_context
+            )
+        ]
+        assert results
+
+        mock_request_handler.on_message_send_stream.assert_called_once()
+        call_context = mock_request_handler.on_message_send_stream.call_args[0][
+            1
+        ]
+        assert isinstance(call_context, ServerCallContext)
+        assert call_context.requested_extensions == {'foo', 'bar'}
+
+        mock_grpc_context.set_trailing_metadata.assert_called_once_with(
+            [(HTTP_EXTENSION_HEADER, 'foo'), (HTTP_EXTENSION_HEADER, 'baz')]
+        )