Skip to content

Commit 7ee1478

Browse files
mindpowerkthota-g
mindpower
authored and
kthota-g
committed
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 4e3f68a commit 7ee1478

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, RequestContext
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,
@@ -59,6 +63,7 @@
5963
)
6064
from a2a.utils.errors import ServerError
6165

66+
6267
MINIMAL_TASK: dict[str, Any] = {
6368
'id': 'task_123',
6469
'contextId': 'session-xyz',
@@ -680,6 +685,265 @@ async def test_on_resubscribe_no_existing_task_error(self) -> None:
680685
self.assertIsInstance(collected_events[0].root, JSONRPCErrorResponse)
681686
assert collected_events[0].root.error == TaskNotFoundError()
682687

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

‎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)