Add tests

Author: kthota-g

tests/server/request_handlers/test_jsonrpc_handler.py

@@ -8,7 +8,7 @@
 import pytest
 import httpx
 
-from a2a.server.agent_execution import AgentExecutor
+from a2a.server.agent_execution import AgentExecutor, RequestContext
 from a2a.server.events import (
     QueueManager,
 )
@@ -55,6 +55,7 @@
     TaskStatusUpdateEvent,
     TextPart,
     UnsupportedOperationError,
+    InternalError,
 )
 from a2a.utils.errors import ServerError
 
@@ -183,7 +184,12 @@ async def test_on_cancel_task_not_found(self) -> None:
         mock_task_store.get.assert_called_once_with('nonexistent_id')
         mock_agent_executor.cancel.assert_not_called()
 
-    async def test_on_message_new_message_success(self) -> None:
+    @patch(
+        'a2a.server.agent_execution.simple_request_context_builder.SimpleRequestContextBuilder.build'
+    )
+    async def test_on_message_new_message_success(
+        self, _mock_builder_build: AsyncMock
+    ) -> None:
         mock_agent_executor = AsyncMock(spec=AgentExecutor)
         mock_task_store = AsyncMock(spec=TaskStore)
         request_handler = DefaultRequestHandler(
@@ -194,6 +200,14 @@ async def test_on_message_new_message_success(self) -> None:
         mock_task_store.get.return_value = mock_task
         mock_agent_executor.execute.return_value = None
 
+        _mock_builder_build.return_value = RequestContext(
+            request=MagicMock(),
+            task_id='task_123',
+            context_id='session-xyz',
+            task=None,
+            related_tasks=None,
+        )
+
         async def streaming_coro():
             yield mock_task
 
@@ -279,15 +293,28 @@ async def streaming_coro():
             assert response.root.error == UnsupportedOperationError()  # type: ignore
             mock_agent_executor.execute.assert_called_once()
 
-    async def test_on_message_stream_new_message_success(self) -> None:
+    @patch(
+        'a2a.server.agent_execution.simple_request_context_builder.SimpleRequestContextBuilder.build'
+    )
+    async def test_on_message_stream_new_message_success(
+        self, _mock_builder_build: AsyncMock
+    ) -> None:
         mock_agent_executor = AsyncMock(spec=AgentExecutor)
         mock_task_store = AsyncMock(spec=TaskStore)
         request_handler = DefaultRequestHandler(
             mock_agent_executor, mock_task_store
         )
-        self.mock_agent_card.capabilities = AgentCapabilities(streaming=True)
 
+        self.mock_agent_card.capabilities = AgentCapabilities(streaming=True)
         handler = JSONRPCHandler(self.mock_agent_card, request_handler)
+        _mock_builder_build.return_value = RequestContext(
+            request=MagicMock(),
+            task_id='task_123',
+            context_id='session-xyz',
+            task=None,
+            related_tasks=None,
+        )
+
         events: list[Any] = [
             Task(**MINIMAL_TASK),
             TaskArtifactUpdateEvent(
@@ -462,8 +489,11 @@ async def test_get_push_notification_success(self) -> None:
         )
         assert get_response.root.result == task_push_config  # type: ignore
 
+    @patch(
+        'a2a.server.agent_execution.simple_request_context_builder.SimpleRequestContextBuilder.build'
+    )
     async def test_on_message_stream_new_message_send_push_notification_success(
-        self,
+        self, _mock_builder_build: AsyncMock
     ) -> None:
         mock_agent_executor = AsyncMock(spec=AgentExecutor)
         mock_task_store = AsyncMock(spec=TaskStore)
@@ -475,6 +505,13 @@ async def test_on_message_stream_new_message_send_push_notification_success(
         self.mock_agent_card.capabilities = AgentCapabilities(
             streaming=True, pushNotifications=True
         )
+        _mock_builder_build.return_value = RequestContext(
+            request=MagicMock(),
+            task_id='task_123',
+            context_id='session-xyz',
+            task=None,
+            related_tasks=None,
+        )
 
         handler = JSONRPCHandler(self.mock_agent_card, request_handler)
         events: list[Any] = [
@@ -642,3 +679,66 @@ async def test_on_resubscribe_no_existing_task_error(self) -> None:
         assert len(collected_events) == 1
         self.assertIsInstance(collected_events[0].root, JSONRPCErrorResponse)
         assert collected_events[0].root.error == TaskNotFoundError()
+
+    async def test_on_message_send_task_id_mismatch(self) -> None:
+        mock_agent_executor = AsyncMock(spec=AgentExecutor)
+        mock_task_store = AsyncMock(spec=TaskStore)
+        request_handler = DefaultRequestHandler(
+            mock_agent_executor, mock_task_store
+        )
+        handler = JSONRPCHandler(self.mock_agent_card, request_handler)
+        mock_task = Task(**MINIMAL_TASK)
+        mock_task_store.get.return_value = mock_task
+        mock_agent_executor.execute.return_value = None
+
+        async def streaming_coro():
+            yield mock_task
+
+        with patch(
+            'a2a.server.request_handlers.default_request_handler.EventConsumer.consume_all',
+            return_value=streaming_coro(),
+        ):
+            request = SendMessageRequest(
+                id='1',
+                params=MessageSendParams(message=Message(**MESSAGE_PAYLOAD)),
+            )
+            response = await handler.on_message_send(request)
+            assert mock_agent_executor.execute.call_count == 1
+            self.assertIsInstance(response.root, JSONRPCErrorResponse)
+            self.assertIsInstance(response.root.error, InternalError)  # type: ignore
+
+    async def test_on_message_stream_task_id_mismatch(self) -> None:
+        mock_agent_executor = AsyncMock(spec=AgentExecutor)
+        mock_task_store = AsyncMock(spec=TaskStore)
+        request_handler = DefaultRequestHandler(
+            mock_agent_executor, mock_task_store
+        )
+
+        self.mock_agent_card.capabilities = AgentCapabilities(streaming=True)
+        handler = JSONRPCHandler(self.mock_agent_card, request_handler)
+        events: list[Any] = [Task(**MINIMAL_TASK)]
+
+        async def streaming_coro():
+            for event in events:
+                yield event
+
+        with patch(
+            'a2a.server.request_handlers.default_request_handler.EventConsumer.consume_all',
+            return_value=streaming_coro(),
+        ):
+            mock_task_store.get.return_value = None
+            mock_agent_executor.execute.return_value = None
+            request = SendStreamingMessageRequest(
+                id='1',
+                params=MessageSendParams(message=Message(**MESSAGE_PAYLOAD)),
+            )
+            response = handler.on_message_send_stream(request)
+            assert isinstance(response, AsyncGenerator)
+            collected_events: list[Any] = []
+            async for event in response:
+                collected_events.append(event)
+            assert len(collected_events) == 1
+            self.assertIsInstance(
+                collected_events[0].root, JSONRPCErrorResponse
+            )
+            self.assertIsInstance(collected_events[0].root.error, InternalError)