Skip to content

Commit 43509d8

Browse files
mikeas1mikeas1
committed
Gemini authored: attempt to restore JSON-RPC client tests
1 parent 2cd2fcd commit 43509d8

2 files changed

Lines changed: 261 additions & 12 deletions

File tree

‎src/a2a/utils/helpers.py‎

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -149,19 +149,18 @@ async def async_wrapper(self: Any, *args, **kwargs) -> Any:
149149
return await function(self, *args, **kwargs)
150150

151151
return async_wrapper
152-
else:
153152

154-
@functools.wraps(function)
155-
def sync_wrapper(self: Any, *args, **kwargs) -> Any:
156-
if not expression(self):
157-
final_message = error_message or str(expression)
158-
logger.error(f'Unsupported Operation: {final_message}')
159-
raise ServerError(
160-
UnsupportedOperationError(message=final_message)
161-
)
162-
return function(self, *args, **kwargs)
153+
@functools.wraps(function)
154+
def sync_wrapper(self: Any, *args, **kwargs) -> Any:
155+
if not expression(self):
156+
final_message = error_message or str(expression)
157+
logger.error(f'Unsupported Operation: {final_message}')
158+
raise ServerError(
159+
UnsupportedOperationError(message=final_message)
160+
)
161+
return function(self, *args, **kwargs)
163162

164-
return sync_wrapper
163+
return sync_wrapper
165164

166165
return decorator
167166

‎tests/client/test_jsonrpc_client.py‎

Lines changed: 251 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -528,4 +528,254 @@ async def test_send_message_client_timeout(
528528
with pytest.raises(A2AClientTimeoutError) as exc_info:
529529
await client.send_message(request=params)
530530

531-
assert 'Request timed out' in str(exc_info.value)
531+
assert 'Client Request timed out' in str(exc_info.value)
532+
533+
@pytest.mark.asyncio
534+
async def test_get_task_success(
535+
self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock
536+
):
537+
client = JsonRpcTransport(
538+
httpx_client=mock_httpx_client, agent_card=mock_agent_card
539+
)
540+
params = TaskQueryParams(id='task-abc')
541+
rpc_response = {
542+
'id': '123',
543+
'jsonrpc': '2.0',
544+
'result': MINIMAL_TASK,
545+
}
546+
with patch.object(
547+
client, '_send_request', new_callable=AsyncMock
548+
) as mock_send_request:
549+
mock_send_request.return_value = rpc_response
550+
response = await client.get_task(request=params)
551+
552+
assert isinstance(response, Task)
553+
assert response.model_dump() == Task.model_validate(
554+
MINIMAL_TASK
555+
).model_dump()
556+
mock_send_request.assert_called_once()
557+
sent_payload = mock_send_request.call_args.args[0]
558+
assert sent_payload['method'] == 'tasks/get'
559+
560+
@pytest.mark.asyncio
561+
async def test_cancel_task_success(
562+
self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock
563+
):
564+
client = JsonRpcTransport(
565+
httpx_client=mock_httpx_client, agent_card=mock_agent_card
566+
)
567+
params = TaskIdParams(id='task-abc')
568+
rpc_response = {
569+
'id': '123',
570+
'jsonrpc': '2.0',
571+
'result': MINIMAL_CANCELLED_TASK,
572+
}
573+
with patch.object(
574+
client, '_send_request', new_callable=AsyncMock
575+
) as mock_send_request:
576+
mock_send_request.return_value = rpc_response
577+
response = await client.cancel_task(request=params)
578+
579+
assert isinstance(response, Task)
580+
assert response.model_dump() == Task.model_validate(
581+
MINIMAL_CANCELLED_TASK
582+
).model_dump()
583+
mock_send_request.assert_called_once()
584+
sent_payload = mock_send_request.call_args.args[0]
585+
assert sent_payload['method'] == 'tasks/cancel'
586+
587+
@pytest.mark.asyncio
588+
async def test_set_task_callback_success(
589+
self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock
590+
):
591+
client = JsonRpcTransport(
592+
httpx_client=mock_httpx_client, agent_card=mock_agent_card
593+
)
594+
params = TaskPushNotificationConfig(
595+
task_id='task-abc',
596+
push_notification_config=PushNotificationConfig(
597+
url='http://callback.com'
598+
),
599+
)
600+
rpc_response = {
601+
'id': '123',
602+
'jsonrpc': '2.0',
603+
'result': params.model_dump(mode='json'),
604+
}
605+
with patch.object(
606+
client, '_send_request', new_callable=AsyncMock
607+
) as mock_send_request:
608+
mock_send_request.return_value = rpc_response
609+
response = await client.set_task_callback(request=params)
610+
611+
assert isinstance(response, TaskPushNotificationConfig)
612+
assert response.model_dump() == params.model_dump()
613+
mock_send_request.assert_called_once()
614+
sent_payload = mock_send_request.call_args.args[0]
615+
assert sent_payload['method'] == 'tasks/pushNotificationConfig/set'
616+
617+
@pytest.mark.asyncio
618+
async def test_get_task_callback_success(
619+
self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock
620+
):
621+
client = JsonRpcTransport(
622+
httpx_client=mock_httpx_client, agent_card=mock_agent_card
623+
)
624+
params = TaskIdParams(id='task-abc')
625+
expected_response = TaskPushNotificationConfig(
626+
task_id='task-abc',
627+
push_notification_config=PushNotificationConfig(
628+
url='http://callback.com'
629+
),
630+
)
631+
rpc_response = {
632+
'id': '123',
633+
'jsonrpc': '2.0',
634+
'result': expected_response.model_dump(mode='json'),
635+
}
636+
with patch.object(
637+
client, '_send_request', new_callable=AsyncMock
638+
) as mock_send_request:
639+
mock_send_request.return_value = rpc_response
640+
response = await client.get_task_callback(request=params)
641+
642+
assert isinstance(response, TaskPushNotificationConfig)
643+
assert response.model_dump() == expected_response.model_dump()
644+
mock_send_request.assert_called_once()
645+
sent_payload = mock_send_request.call_args.args[0]
646+
assert sent_payload['method'] == 'tasks/pushNotificationConfig/get'
647+
648+
@pytest.mark.asyncio
649+
@patch('a2a.client.transports.jsonrpc.aconnect_sse')
650+
async def test_send_message_streaming_sse_error(
651+
self,
652+
mock_aconnect_sse: AsyncMock,
653+
mock_httpx_client: AsyncMock,
654+
mock_agent_card: MagicMock,
655+
):
656+
client = JsonRpcTransport(
657+
httpx_client=mock_httpx_client, agent_card=mock_agent_card
658+
)
659+
params = MessageSendParams(
660+
message=create_text_message_object(content='Hello stream')
661+
)
662+
mock_event_source = AsyncMock(spec=EventSource)
663+
mock_event_source.aiter_sse.side_effect = SSEError(
664+
'Simulated SSE error'
665+
)
666+
mock_aconnect_sse.return_value.__aenter__.return_value = (
667+
mock_event_source
668+
)
669+
670+
with pytest.raises(A2AClientHTTPError):
671+
_ = [
672+
item
673+
async for item in client.send_message_streaming(request=params)
674+
]
675+
676+
@pytest.mark.asyncio
677+
@patch('a2a.client.transports.jsonrpc.aconnect_sse')
678+
async def test_send_message_streaming_json_error(
679+
self,
680+
mock_aconnect_sse: AsyncMock,
681+
mock_httpx_client: AsyncMock,
682+
mock_agent_card: MagicMock,
683+
):
684+
client = JsonRpcTransport(
685+
httpx_client=mock_httpx_client, agent_card=mock_agent_card
686+
)
687+
params = MessageSendParams(
688+
message=create_text_message_object(content='Hello stream')
689+
)
690+
sse_event = ServerSentEvent(data='{invalid json')
691+
mock_event_source = AsyncMock(spec=EventSource)
692+
mock_event_source.aiter_sse.return_value = async_iterable_from_list(
693+
[sse_event]
694+
)
695+
mock_aconnect_sse.return_value.__aenter__.return_value = (
696+
mock_event_source
697+
)
698+
699+
with pytest.raises(A2AClientJSONError):
700+
_ = [
701+
item
702+
async for item in client.send_message_streaming(request=params)
703+
]
704+
705+
@pytest.mark.asyncio
706+
@patch('a2a.client.transports.jsonrpc.aconnect_sse')
707+
async def test_send_message_streaming_request_error(
708+
self,
709+
mock_aconnect_sse: AsyncMock,
710+
mock_httpx_client: AsyncMock,
711+
mock_agent_card: MagicMock,
712+
):
713+
client = JsonRpcTransport(
714+
httpx_client=mock_httpx_client, agent_card=mock_agent_card
715+
)
716+
params = MessageSendParams(
717+
message=create_text_message_object(content='Hello stream')
718+
)
719+
mock_event_source = AsyncMock(spec=EventSource)
720+
mock_event_source.aiter_sse.side_effect = httpx.RequestError(
721+
'Simulated request error', request=MagicMock()
722+
)
723+
mock_aconnect_sse.return_value.__aenter__.return_value = (
724+
mock_event_source
725+
)
726+
727+
with pytest.raises(A2AClientHTTPError):
728+
_ = [
729+
item
730+
async for item in client.send_message_streaming(request=params)
731+
]
732+
733+
@pytest.mark.asyncio
734+
async def test_get_card_no_card_provided(self, mock_httpx_client: AsyncMock):
735+
client = JsonRpcTransport(
736+
httpx_client=mock_httpx_client, url=self.AGENT_URL
737+
)
738+
mock_response = AsyncMock(spec=httpx.Response)
739+
mock_response.status_code = 200
740+
mock_response.json.return_value = AGENT_CARD.model_dump(mode='json')
741+
mock_httpx_client.get.return_value = mock_response
742+
743+
card = await client.get_card()
744+
745+
assert card == AGENT_CARD
746+
mock_httpx_client.get.assert_called_once()
747+
748+
@pytest.mark.asyncio
749+
async def test_get_card_with_extended_card_support(
750+
self, mock_httpx_client: AsyncMock
751+
):
752+
agent_card = AGENT_CARD.model_copy(
753+
update={'supports_authenticated_extended_card': True}
754+
)
755+
client = JsonRpcTransport(
756+
httpx_client=mock_httpx_client, agent_card=agent_card
757+
)
758+
759+
rpc_response = {
760+
'id': '123',
761+
'jsonrpc': '2.0',
762+
'result': AGENT_CARD_EXTENDED.model_dump(mode='json'),
763+
}
764+
with patch.object(
765+
client, '_send_request', new_callable=AsyncMock
766+
) as mock_send_request:
767+
mock_send_request.return_value = rpc_response
768+
card = await client.get_card()
769+
770+
assert card == agent_card
771+
mock_send_request.assert_called_once()
772+
sent_payload = mock_send_request.call_args.args[0]
773+
assert sent_payload['method'] == 'agent/getAuthenticatedExtendedCard'
774+
775+
@pytest.mark.asyncio
776+
async def test_close(self, mock_httpx_client: AsyncMock):
777+
client = JsonRpcTransport(
778+
httpx_client=mock_httpx_client, url=self.AGENT_URL
779+
)
780+
await client.close()
781+
mock_httpx_client.aclose.assert_called_once()

0 commit comments

Comments
 (0)