Add tests

Author: ishymko

src/a2a/client/transports/grpc.py

@@ -4,9 +4,8 @@
 from functools import wraps
 from typing import Any, NoReturn
 
-import a2a.utils.errors
-
 from a2a.client.errors import A2AClientError
+from a2a.utils.errors import JSON_RPC_ERROR_CODE_MAP
 
 
 try:
@@ -48,18 +47,18 @@
 
 logger = logging.getLogger(__name__)
 
+_A2A_ERROR_NAME_TO_CLS = {
+    error_type.__name__: error_type for error_type in JSON_RPC_ERROR_CODE_MAP
+}
+
 
 def _map_grpc_error(e: grpc.aio.AioRpcError) -> NoReturn:
     details = e.details()
     if isinstance(details, str) and ': ' in details:
         error_type_name, error_message = details.split(': ', 1)
         # TODO(#723): Resolving imports by name is a temporary hack until proper error handling structure is added in #723.
-        exception_cls = getattr(a2a.utils.errors, error_type_name, None)
-        if (
-            exception_cls
-            and isinstance(exception_cls, type)
-            and issubclass(exception_cls, a2a.utils.errors.A2AError)
-        ):
+        exception_cls = _A2A_ERROR_NAME_TO_CLS.get(error_type_name)
+        if exception_cls:
             raise exception_cls(error_message) from e
     raise A2AClientError(f'gRPC Error {e.code().name}: {e.details()}') from e
 

src/a2a/client/transports/rest.py

@@ -10,8 +10,6 @@
 from google.protobuf.message import Message
 from httpx_sse import SSEError, aconnect_sse
 
-import a2a.utils.errors
-
 from a2a.client.errors import A2AClientError
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
 from a2a.client.transports.base import ClientTransport
@@ -34,12 +32,16 @@
     Task,
     TaskPushNotificationConfig,
 )
-from a2a.utils.errors import MethodNotFoundError
+from a2a.utils.errors import JSON_RPC_ERROR_CODE_MAP, MethodNotFoundError
 from a2a.utils.telemetry import SpanKind, trace_class
 
 
 logger = logging.getLogger(__name__)
 
+_A2A_ERROR_NAME_TO_CLS = {
+    error_type.__name__: error_type for error_type in JSON_RPC_ERROR_CODE_MAP
+}
+
 
 @trace_class(kind=SpanKind.CLIENT)
 class RestTransport(ClientTransport):
@@ -103,12 +105,9 @@ def _handle_http_error(self, e: httpx.HTTPStatusError) -> NoReturn:
             message = error_data.get('message', str(e))
 
             if isinstance(error_type, str):
-                exception_cls = getattr(a2a.utils.errors, error_type, None)
-                if (
-                    exception_cls
-                    and isinstance(exception_cls, type)
-                    and issubclass(exception_cls, a2a.utils.errors.A2AError)
-                ):
+                # TODO(#723): Resolving imports by name is a temporary hack until proper error handling structure is added in #723.
+                exception_cls = _A2A_ERROR_NAME_TO_CLS.get(error_type)
+                if exception_cls:
                     raise exception_cls(message) from e
         except (json.JSONDecodeError, ValueError):
             pass

tests/client/transports/test_grpc_client.py

@@ -5,19 +5,18 @@
 
 from a2a.client.transports.grpc import GrpcTransport
 from a2a.extensions.common import HTTP_EXTENSION_HEADER
-from a2a.types import a2a_pb2, a2a_pb2_grpc
+from a2a.types import a2a_pb2
 from a2a.types.a2a_pb2 import (
     AgentCapabilities,
-    AgentInterface,
     AgentCard,
+    AgentInterface,
     Artifact,
     AuthenticationInfo,
     CreateTaskPushNotificationConfigRequest,
     DeleteTaskPushNotificationConfigRequest,
     GetTaskPushNotificationConfigRequest,
-    ListTaskPushNotificationConfigsRequest,
-    ListTaskPushNotificationConfigsResponse,
     GetTaskRequest,
+    ListTaskPushNotificationConfigsRequest,
     Message,
     Part,
     PushNotificationConfig,
@@ -30,7 +29,8 @@
     TaskStatus,
     TaskStatusUpdateEvent,
 )
-from a2a.utils import get_text_parts, proto_utils
+from a2a.utils import get_text_parts
+from a2a.utils.errors import JSON_RPC_ERROR_CODE_MAP
 
 
 @pytest.fixture
@@ -226,6 +226,29 @@ async def test_send_message_task_response(
     assert response.task.id == sample_task.id
 
 
+@pytest.mark.parametrize('error_cls', list(JSON_RPC_ERROR_CODE_MAP.keys()))
+@pytest.mark.asyncio
+async def test_grpc_mapped_errors(
+    grpc_transport: GrpcTransport,
+    mock_grpc_stub: AsyncMock,
+    sample_message_send_params: SendMessageRequest,
+    error_cls,
+) -> None:
+    """Test handling of mapped gRPC error responses."""
+    error_details = f'{error_cls.__name__}: Mapped Error'
+
+    # We must trigger it from a standard transport method call, for example `send_message`.
+    mock_grpc_stub.SendMessage.side_effect = grpc.aio.AioRpcError(
+        code=grpc.StatusCode.INTERNAL,
+        initial_metadata=grpc.aio.Metadata(),
+        trailing_metadata=grpc.aio.Metadata(),
+        details=error_details,
+    )
+
+    with pytest.raises(error_cls):
+        await grpc_transport.send_message(sample_message_send_params)
+
+
 @pytest.mark.asyncio
 async def test_send_message_message_response(
     grpc_transport: GrpcTransport,

tests/client/transports/test_jsonrpc_client.py

@@ -1,30 +1,27 @@
 """Tests for the JSON-RPC client transport."""
 
 import json
-from google.protobuf import json_format
-from unittest import mock
+
 from unittest.mock import AsyncMock, MagicMock, patch
 from uuid import uuid4
 
 import httpx
 import pytest
-import respx
+
+from google.protobuf import json_format
 from httpx_sse import EventSource, SSEError
 
 from a2a.client.errors import A2AClientError
-from a2a.utils.errors import InvalidRequestError
 from a2a.client.transports.jsonrpc import JsonRpcTransport
 from a2a.types.a2a_pb2 import (
     AgentCapabilities,
-    AgentInterface,
     AgentCard,
+    AgentInterface,
     CancelTaskRequest,
-    CreateTaskPushNotificationConfigRequest,
     DeleteTaskPushNotificationConfigRequest,
     GetTaskPushNotificationConfigRequest,
-    ListTaskPushNotificationConfigsRequest,
-    ListTaskPushNotificationConfigsResponse,
     GetTaskRequest,
+    ListTaskPushNotificationConfigsRequest,
     Message,
     Part,
     SendMessageConfiguration,
@@ -33,8 +30,8 @@
     Task,
     TaskPushNotificationConfig,
     TaskState,
-    TaskStatus,
 )
+from a2a.utils.errors import JSON_RPC_ERROR_CODE_MAP
 
 
 @pytest.fixture
@@ -174,16 +171,19 @@ async def test_send_message_success(self, transport, mock_httpx_client):
         payload = call_args[1]['json']
         assert payload['method'] == 'SendMessage'
 
+    @pytest.mark.parametrize(
+        'error_cls, error_code', JSON_RPC_ERROR_CODE_MAP.items()
+    )
     @pytest.mark.asyncio
     async def test_send_message_jsonrpc_error(
-        self, transport, mock_httpx_client
+        self, transport, mock_httpx_client, error_cls, error_code
     ):
-        """Test handling of JSON-RPC error response."""
+        """Test handling of JSON-RPC mapped error response."""
         mock_response = MagicMock()
         mock_response.json.return_value = {
             'jsonrpc': '2.0',
             'id': '1',
-            'error': {'code': -32600, 'message': 'Invalid Request'},
+            'error': {'code': error_code, 'message': 'Mapped Error'},
             'result': None,
         }
         mock_response.raise_for_status = MagicMock()
@@ -192,7 +192,7 @@ async def test_send_message_jsonrpc_error(
         request = create_send_message_request()
 
         # The transport raises the specific A2AError mapped from code
-        with pytest.raises(InvalidRequestError):
+        with pytest.raises(error_cls):
             await transport.send_message(request)
 
     @pytest.mark.asyncio

tests/client/transports/test_rest_client.py

@@ -17,11 +17,10 @@
     AgentInterface,
     DeleteTaskPushNotificationConfigRequest,
     ListTaskPushNotificationConfigsRequest,
-    ListTaskPushNotificationConfigsResponse,
     SendMessageRequest,
-    TaskPushNotificationConfig,
 )
 from a2a.utils.constants import TransportProtocol
+from a2a.utils.errors import JSON_RPC_ERROR_CODE_MAP
 
 
 @pytest.fixture
@@ -95,6 +94,42 @@ async def test_send_message_streaming_timeout(
         assert 'Client Request timed out' in str(exc_info.value)
 
 
+    @pytest.mark.parametrize('error_cls', list(JSON_RPC_ERROR_CODE_MAP.keys()))
+    @pytest.mark.asyncio
+    async def test_rest_mapped_errors(
+        self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock, error_cls
+    ):
+        """Test handling of mapped REST HTTP error responses."""
+        client = RestTransport(
+            httpx_client=mock_httpx_client,
+            agent_card=mock_agent_card,
+            url='http://agent.example.com/api',
+        )
+        params = SendMessageRequest(
+            message=create_text_message_object(content='Hello')
+        )
+
+        mock_build_request = MagicMock(
+            return_value=AsyncMock(spec=httpx.Request)
+        )
+        mock_httpx_client.build_request = mock_build_request
+
+        mock_response = AsyncMock(spec=httpx.Response)
+        mock_response.status_code = 500
+        mock_response.json.return_value = {'type': error_cls.__name__, 'message': 'Mapped Error'}
+
+        error = httpx.HTTPStatusError(
+            'Server Error',
+            request=httpx.Request('POST', 'http://test.url'),
+            response=mock_response,
+        )
+
+        mock_httpx_client.send.side_effect = error
+
+        with pytest.raises(error_cls):
+            await client.send_message(request=params)
+
+
 class TestRestTransportExtensions:
     @pytest.mark.asyncio
     async def test_send_message_with_default_extensions(