Skip to content

Commit 3b811ca

Browse files
mindpowermindpower
authored
test: Adding 8 tests for jsonrpc_handler.py and also fix minor waring in test_integration.py (#72)
* test: Adding 8 tests for jsonrpc_handler.py and also fix minor waring in test_integration.py * test: remove comments
1 parent 4ca8eff commit 3b811ca

File tree

2 files changed

+269
-5
lines changed

2 files changed

+269
-5
lines changed

‎tests/server/request_handlers/test_jsonrpc_handler.py‎

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