Skip to content

Commit 89ab174

Browse files
mindpowermindpower
committed
test: Adding 8 tests for jsonrpc_handler.py and also fix minor waring in test_integration.py
1 parent 4ca8eff commit 89ab174

File tree

2 files changed

+271
-5
lines changed

2 files changed

+271
-5
lines changed

‎tests/server/request_handlers/test_jsonrpc_handler.py‎

Lines changed: 269 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,15 @@
33

44
from collections.abc import AsyncGenerator
55
from typing import Any
6-
from unittest.mock import AsyncMock, MagicMock, patch, call
6+
from unittest.mock import AsyncMock, MagicMock, call, patch
77

8-
import pytest
98
import httpx
9+
import pytest
1010

1111
from a2a.server.agent_execution import AgentExecutor
12+
from a2a.server.agent_execution.request_context_builder import (
13+
RequestContextBuilder,
14+
)
1215
from a2a.server.events import (
1316
QueueManager,
1417
)
@@ -30,14 +33,15 @@
3033
GetTaskRequest,
3134
GetTaskResponse,
3235
GetTaskSuccessResponse,
36+
InternalError,
3337
JSONRPCErrorResponse,
3438
Message,
3539
MessageSendConfiguration,
3640
MessageSendParams,
3741
Part,
42+
PushNotificationConfig,
3843
SendMessageRequest,
3944
SendMessageSuccessResponse,
40-
PushNotificationConfig,
4145
SendStreamingMessageRequest,
4246
SendStreamingMessageSuccessResponse,
4347
SetTaskPushNotificationConfigRequest,
@@ -58,6 +62,7 @@
5862
)
5963
from a2a.utils.errors import ServerError
6064

65+
6166
MINIMAL_TASK: dict[str, Any] = {
6267
'id': 'task_123',
6368
'contextId': 'session-xyz',
@@ -642,3 +647,264 @@ async def test_on_resubscribe_no_existing_task_error(self) -> None:
642647
assert len(collected_events) == 1
643648
self.assertIsInstance(collected_events[0].root, JSONRPCErrorResponse)
644649
assert collected_events[0].root.error == TaskNotFoundError()
650+
651+
# Additional tests to add to your TestJSONRPCtHandler class in test_jsonrpc_handler.py
652+
653+
async def test_streaming_not_supported_error(
654+
self,
655+
) -> None:
656+
"""Test that on_message_send_stream raises an error when streaming not supported."""
657+
# Arrange
658+
mock_agent_executor = AsyncMock(spec=AgentExecutor)
659+
mock_task_store = AsyncMock(spec=TaskStore)
660+
request_handler = DefaultRequestHandler(
661+
mock_agent_executor, mock_task_store
662+
)
663+
# Create agent card with streaming capability disabled
664+
self.mock_agent_card.capabilities = AgentCapabilities(streaming=False)
665+
handler = JSONRPCHandler(self.mock_agent_card, request_handler)
666+
667+
# Act & Assert
668+
request = SendStreamingMessageRequest(
669+
id='1',
670+
params=MessageSendParams(message=Message(**MESSAGE_PAYLOAD)),
671+
)
672+
673+
# Should raise ServerError about streaming not supported
674+
with self.assertRaises(ServerError) as context:
675+
async for _ in handler.on_message_send_stream(request):
676+
pass
677+
678+
aaa = context.exception
679+
self.assertEqual(
680+
str(context.exception.error.message),
681+
'Streaming is not supported by the agent',
682+
)
683+
684+
async def test_push_notifications_not_supported_error(self) -> None:
685+
"""Test that set_push_notification raises an error when push notifications not supported."""
686+
# Arrange
687+
mock_agent_executor = AsyncMock(spec=AgentExecutor)
688+
mock_task_store = AsyncMock(spec=TaskStore)
689+
request_handler = DefaultRequestHandler(
690+
mock_agent_executor, mock_task_store
691+
)
692+
# Create agent card with push notifications capability disabled
693+
self.mock_agent_card.capabilities = AgentCapabilities(
694+
pushNotifications=False, streaming=True
695+
)
696+
handler = JSONRPCHandler(self.mock_agent_card, request_handler)
697+
698+
# Act & Assert
699+
task_push_config = TaskPushNotificationConfig(
700+
taskId='task_123',
701+
pushNotificationConfig=PushNotificationConfig(
702+
url='http://example.com'
703+
),
704+
)
705+
request = SetTaskPushNotificationConfigRequest(
706+
id='1', params=task_push_config
707+
)
708+
709+
# Should raise ServerError about push notifications not supported
710+
with self.assertRaises(ServerError) as context:
711+
await handler.set_push_notification(request)
712+
713+
self.assertEqual(
714+
str(context.exception.error.message),
715+
'Push notifications are not supported by the agent',
716+
)
717+
718+
async def test_on_get_push_notification_no_push_notifier(self) -> None:
719+
"""Test get_push_notification with no push notifier configured."""
720+
# Arrange
721+
mock_agent_executor = AsyncMock(spec=AgentExecutor)
722+
mock_task_store = AsyncMock(spec=TaskStore)
723+
# Create request handler without a push notifier
724+
request_handler = DefaultRequestHandler(
725+
mock_agent_executor, mock_task_store
726+
)
727+
self.mock_agent_card.capabilities = AgentCapabilities(
728+
pushNotifications=True
729+
)
730+
handler = JSONRPCHandler(self.mock_agent_card, request_handler)
731+
732+
mock_task = Task(**MINIMAL_TASK)
733+
mock_task_store.get.return_value = mock_task
734+
735+
# Act
736+
get_request = GetTaskPushNotificationConfigRequest(
737+
id='1', params=TaskIdParams(id=mock_task.id)
738+
)
739+
response = await handler.get_push_notification(get_request)
740+
741+
# Assert
742+
self.assertIsInstance(response.root, JSONRPCErrorResponse)
743+
self.assertEqual(response.root.error, UnsupportedOperationError())
744+
745+
async def test_on_set_push_notification_no_push_notifier(self) -> None:
746+
"""Test set_push_notification with no push notifier configured."""
747+
# Arrange
748+
mock_agent_executor = AsyncMock(spec=AgentExecutor)
749+
mock_task_store = AsyncMock(spec=TaskStore)
750+
# Create request handler without a push notifier
751+
request_handler = DefaultRequestHandler(
752+
mock_agent_executor, mock_task_store
753+
)
754+
self.mock_agent_card.capabilities = AgentCapabilities(
755+
pushNotifications=True
756+
)
757+
handler = JSONRPCHandler(self.mock_agent_card, request_handler)
758+
759+
mock_task = Task(**MINIMAL_TASK)
760+
mock_task_store.get.return_value = mock_task
761+
762+
# Act
763+
task_push_config = TaskPushNotificationConfig(
764+
taskId=mock_task.id,
765+
pushNotificationConfig=PushNotificationConfig(
766+
url='http://example.com'
767+
),
768+
)
769+
request = SetTaskPushNotificationConfigRequest(
770+
id='1', params=task_push_config
771+
)
772+
response = await handler.set_push_notification(request)
773+
774+
# Assert
775+
self.assertIsInstance(response.root, JSONRPCErrorResponse)
776+
self.assertEqual(response.root.error, UnsupportedOperationError())
777+
778+
async def test_on_message_send_internal_error(self) -> None:
779+
"""Test on_message_send with an internal error."""
780+
# Arrange
781+
mock_agent_executor = AsyncMock(spec=AgentExecutor)
782+
mock_task_store = AsyncMock(spec=TaskStore)
783+
request_handler = DefaultRequestHandler(
784+
mock_agent_executor, mock_task_store
785+
)
786+
handler = JSONRPCHandler(self.mock_agent_card, request_handler)
787+
788+
# Make the request handler raise an Internal error without specifying an error type
789+
async def raise_server_error(*args, **kwargs):
790+
raise ServerError(InternalError(message='Internal Error'))
791+
792+
# Patch the method to raise an error
793+
with patch.object(
794+
request_handler, 'on_message_send', side_effect=raise_server_error
795+
):
796+
# Act
797+
request = SendMessageRequest(
798+
id='1',
799+
params=MessageSendParams(message=Message(**MESSAGE_PAYLOAD)),
800+
)
801+
response = await handler.on_message_send(request)
802+
803+
# Assert
804+
self.assertIsInstance(response.root, JSONRPCErrorResponse)
805+
self.assertIsInstance(response.root.error, InternalError)
806+
807+
async def test_on_message_stream_internal_error(self) -> None:
808+
"""Test on_message_send_stream with an internal error."""
809+
# Arrange
810+
mock_agent_executor = AsyncMock(spec=AgentExecutor)
811+
mock_task_store = AsyncMock(spec=TaskStore)
812+
request_handler = DefaultRequestHandler(
813+
mock_agent_executor, mock_task_store
814+
)
815+
self.mock_agent_card.capabilities = AgentCapabilities(streaming=True)
816+
handler = JSONRPCHandler(self.mock_agent_card, request_handler)
817+
818+
# Make the request handler raise an Internal error without specifying an error type
819+
async def raise_server_error(*args, **kwargs):
820+
raise ServerError(InternalError(message='Internal Error'))
821+
yield # Need this to make it an async generator
822+
823+
# Patch the method to raise an error
824+
with patch.object(
825+
request_handler,
826+
'on_message_send_stream',
827+
return_value=raise_server_error(),
828+
):
829+
# Act
830+
request = SendStreamingMessageRequest(
831+
id='1',
832+
params=MessageSendParams(message=Message(**MESSAGE_PAYLOAD)),
833+
)
834+
835+
# Get the single error response
836+
responses = []
837+
async for response in handler.on_message_send_stream(request):
838+
responses.append(response)
839+
840+
# Assert
841+
self.assertEqual(len(responses), 1)
842+
self.assertIsInstance(responses[0].root, JSONRPCErrorResponse)
843+
self.assertIsInstance(responses[0].root.error, InternalError)
844+
845+
async def test_default_request_handler_with_custom_components(self) -> None:
846+
"""Test DefaultRequestHandler initialization with custom components."""
847+
# Arrange
848+
mock_agent_executor = AsyncMock(spec=AgentExecutor)
849+
mock_task_store = AsyncMock(spec=TaskStore)
850+
mock_queue_manager = AsyncMock(spec=QueueManager)
851+
mock_push_notifier = AsyncMock(spec=PushNotifier)
852+
mock_request_context_builder = AsyncMock(spec=RequestContextBuilder)
853+
854+
# Act
855+
handler = DefaultRequestHandler(
856+
agent_executor=mock_agent_executor,
857+
task_store=mock_task_store,
858+
queue_manager=mock_queue_manager,
859+
push_notifier=mock_push_notifier,
860+
request_context_builder=mock_request_context_builder,
861+
)
862+
863+
# Assert
864+
self.assertEqual(handler.agent_executor, mock_agent_executor)
865+
self.assertEqual(handler.task_store, mock_task_store)
866+
self.assertEqual(handler._queue_manager, mock_queue_manager)
867+
self.assertEqual(handler._push_notifier, mock_push_notifier)
868+
self.assertEqual(
869+
handler._request_context_builder, mock_request_context_builder
870+
)
871+
872+
async def test_on_message_send_error_handling(self) -> None:
873+
"""Test error handling in on_message_send when consuming raises ServerError."""
874+
# Arrange
875+
mock_agent_executor = AsyncMock(spec=AgentExecutor)
876+
mock_task_store = AsyncMock(spec=TaskStore)
877+
request_handler = DefaultRequestHandler(
878+
mock_agent_executor, mock_task_store
879+
)
880+
handler = JSONRPCHandler(self.mock_agent_card, request_handler)
881+
882+
# Let task exist
883+
mock_task = Task(**MINIMAL_TASK)
884+
mock_task_store.get.return_value = mock_task
885+
886+
# Set up consume_and_break_on_interrupt to raise ServerError
887+
async def consume_raises_error(*args, **kwargs):
888+
raise ServerError(error=UnsupportedOperationError())
889+
890+
with patch(
891+
'a2a.server.tasks.result_aggregator.ResultAggregator.consume_and_break_on_interrupt',
892+
side_effect=consume_raises_error,
893+
):
894+
# Act
895+
request = SendMessageRequest(
896+
id='1',
897+
params=MessageSendParams(
898+
message=Message(
899+
**MESSAGE_PAYLOAD,
900+
taskId=mock_task.id,
901+
contextId=mock_task.contextId,
902+
)
903+
),
904+
)
905+
906+
response = await handler.on_message_send(request)
907+
908+
# Assert
909+
self.assertIsInstance(response.root, JSONRPCErrorResponse)
910+
self.assertEqual(response.root.error, UnsupportedOperationError())

‎tests/server/test_integration.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -230,7 +230,7 @@ def test_cancel_task(client: TestClient, handler: mock.AsyncMock):
230230
"""Test cancelling a task."""
231231
# Setup mock response
232232
task_status = TaskStatus(**MINIMAL_TASK_STATUS)
233-
task_status.state = TaskState.canceled # 'cancelled' #
233+
task_status.state = TaskState.canceled # 'cancelled' #
234234
task = Task(
235235
id='task1', contextId='ctx1', state='cancelled', status=task_status
236236
)
@@ -543,7 +543,7 @@ async def stream_generator():
543543

544544
def test_invalid_json(client: TestClient):
545545
"""Test handling invalid JSON."""
546-
response = client.post('/', data='This is not JSON')
546+
response = client.post('/', content=b'This is not JSON') # Use bytes
547547
assert response.status_code == 200 # JSON-RPC errors still return 200
548548
data = response.json()
549549
assert 'error' in data

0 commit comments

Comments
 (0)