Fix broken tests

Author: mikeas1

src/a2a/client/transports/rest.py

@@ -310,24 +310,14 @@ async def resubscribe(
         Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent | Message
     ]:
         """Reconnects to get task updates."""
-        pb = a2a_pb2.TaskSubscriptionRequest(
-            name=f'tasks/{request.id}',
-        )
-        payload = MessageToDict(pb)
-        payload, modified_kwargs = await self._apply_interceptors(
-            payload,
-            self._get_http_args(context),
-            context,
-        )
-
-        modified_kwargs.setdefault('timeout', None)
+        http_kwargs = self._get_http_args(context) or {}
+        http_kwargs.setdefault('timeout', None)
 
         async with aconnect_sse(
             self.httpx_client,
             'GET',
             f'{self.url}/v1/tasks/{request.id}:subscribe',
-            json=payload,
-            **modified_kwargs,
+            **http_kwargs,
         ) as event_source:
             try:
                 async for sse in event_source.aiter_sse():

src/a2a/server/apps/rest/rest_app.py

@@ -15,10 +15,7 @@
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.request_handler import RequestHandler
 from a2a.server.request_handlers.rest_handler import RESTHandler
-from a2a.types import (
-    AgentCard,
-    AuthenticatedExtendedCardNotConfiguredError,
-)
+from a2a.types import AgentCard, AuthenticatedExtendedCardNotConfiguredError
 from a2a.utils.error_handlers import (
     rest_error_handler,
     rest_stream_error_handler,
@@ -61,21 +58,7 @@ def __init__(
     @rest_error_handler
     async def _handle_request(
         self,
-        method: Callable[
-            [Request, ServerCallContext], Awaitable[dict[str, Any]]
-        ],
-        request: Request,
-    ) -> Response:
-        call_context = self._context_builder.build(request)
-        response = await method(request, call_context)
-        return JSONResponse(content=response)
-
-    @rest_error_handler
-    async def _handle_list_request(
-        self,
-        method: Callable[
-            [Request, ServerCallContext], Awaitable[list[dict[str, Any]]]
-        ],
+        method: Callable[[Request, ServerCallContext], Awaitable[Any]],
         request: Request,
     ) -> Response:
         call_context = self._context_builder.build(request)
@@ -85,15 +68,13 @@ async def _handle_list_request(
     @rest_stream_error_handler
     async def _handle_streaming_request(
         self,
-        method: Callable[
-            [Request, ServerCallContext], AsyncIterable[dict[str, Any]]
-        ],
+        method: Callable[[Request, ServerCallContext], AsyncIterable[Any]],
         request: Request,
     ) -> EventSourceResponse:
         call_context = self._context_builder.build(request)
 
         async def event_generator(
-            stream: AsyncIterable[dict[str, Any]],
+            stream: AsyncIterable[Any],
         ) -> AsyncIterator[dict[str, dict[str, Any]]]:
             async for item in stream:
                 yield {'data': item}
@@ -188,10 +169,10 @@ def routes(self) -> dict[tuple[str, str], Callable[[Request], Any]]:
                 '/v1/tasks/{id}/pushNotificationConfigs',
                 'GET',
             ): functools.partial(
-                self._handle_list_request, self.handler.list_push_notifications
+                self._handle_request, self.handler.list_push_notifications
             ),
             ('/v1/tasks', 'GET'): functools.partial(
-                self._handle_list_request, self.handler.list_tasks
+                self._handle_request, self.handler.list_tasks
             ),
         }
         if self.agent_card.supports_authenticated_extended_card:

src/a2a/server/request_handlers/rest_handler.py

@@ -3,7 +3,7 @@
 from collections.abc import AsyncIterable, AsyncIterator
 from typing import Any
 
-from google.protobuf.json_format import MessageToDict, Parse
+from google.protobuf.json_format import MessageToDict, MessageToJson, Parse
 from starlette.requests import Request
 
 from a2a.grpc import a2a_pb2
@@ -86,7 +86,7 @@ async def on_message_send_stream(
         self,
         request: Request,
         context: ServerCallContext,
-    ) -> AsyncIterator[dict[str, Any]]:
+    ) -> AsyncIterator[str]:
         """Handles the 'message/stream' REST method.
 
         Yields response objects as they are produced by the underlying handler's stream.
@@ -96,7 +96,7 @@ async def on_message_send_stream(
             context: Context provided by the server.
 
         Yields:
-            `dict` objects containing streaming events
+            JSON serialized objects containing streaming events
             (Task, Message, TaskStatusUpdateEvent, TaskArtifactUpdateEvent) as JSON
         """
         body = await request.body()
@@ -110,7 +110,7 @@ async def on_message_send_stream(
             a2a_request, context
         ):
             response = proto_utils.ToProto.stream_response(event)
-            yield MessageToDict(response)
+            yield MessageToJson(response)
 
     async def on_cancel_task(
         self,
@@ -142,7 +142,7 @@ async def on_resubscribe_to_task(
         self,
         request: Request,
         context: ServerCallContext,
-    ) -> AsyncIterable[dict[str, Any]]:
+    ) -> AsyncIterable[str]:
         """Handles the 'tasks/resubscribe' REST method.
 
         Yields response objects as they are produced by the underlying handler's stream.
@@ -152,13 +152,13 @@ async def on_resubscribe_to_task(
             context: Context provided by the server.
 
         Yields:
-            `dict` containing streaming events
+            JSON serialized objects containing streaming events
         """
         task_id = request.path_params['id']
         async for event in self.request_handler.on_resubscribe_to_task(
             TaskIdParams(id=task_id), context
         ):
-            yield (MessageToDict(proto_utils.ToProto.stream_response(event)))
+            yield MessageToJson(proto_utils.ToProto.stream_response(event))
 
     async def get_push_notification(
         self,
@@ -262,7 +262,7 @@ async def list_push_notifications(
         self,
         request: Request,
         context: ServerCallContext,
-    ) -> list[dict[str, Any]]:
+    ) -> dict[str, Any]:
         """Handles the 'tasks/pushNotificationConfig/list' REST method.
 
         This method is currently not implemented.
@@ -283,7 +283,7 @@ async def list_tasks(
         self,
         request: Request,
         context: ServerCallContext,
-    ) -> list[dict[str, Any]]:
+    ) -> dict[str, Any]:
         """Handles the 'tasks/list' REST method.
 
         This method is currently not implemented.

tests/client/test_jsonrpc_client.py

@@ -334,15 +334,15 @@ def test_init_with_url(self, mock_httpx_client: AsyncMock):
         assert client.url == self.AGENT_URL
         assert client.httpx_client == mock_httpx_client
 
-    def test_init_with_agent_card_and_url_prioritizes_agent_card(
+    def test_init_with_agent_card_and_url_prioritizes_url(
         self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock
     ):
         client = JsonRpcTransport(
             httpx_client=mock_httpx_client,
             agent_card=mock_agent_card,
             url='http://otherurl.com',
         )
-        assert client.url == mock_agent_card.url
+        assert client.url == 'http://otherurl.com'
 
     def test_init_raises_value_error_if_no_card_or_url(
         self, mock_httpx_client: AsyncMock

tests/client/test_legacy_client.py

@@ -1,83 +1,10 @@
 """Tests for the legacy client compatibility layer."""
 
-from unittest.mock import AsyncMock, MagicMock
-
-import httpx
-import pytest
-
-from a2a.client import A2AClient, A2AGrpcClient
-from a2a.types import (
-    AgentCapabilities,
-    AgentCard,
-    Message,
-    MessageSendParams,
-    Part,
-    Role,
-    SendMessageRequest,
-    Task,
-    TaskQueryParams,
-    TaskState,
-    TaskStatus,
-    TextPart,
-)
-
-
-@pytest.fixture
-def mock_httpx_client() -> AsyncMock:
-    return AsyncMock(spec=httpx.AsyncClient)
-
-
-@pytest.fixture
-def mock_grpc_stub() -> AsyncMock:
-    stub = AsyncMock()
-    stub._channel = MagicMock()
-    return stub
-
-
-@pytest.fixture
-def jsonrpc_agent_card() -> AgentCard:
-    return AgentCard(
-        name='Test Agent',
-        description='A test agent',
-        url='http://test.agent.com/rpc',
-        version='1.0.0',
-        capabilities=AgentCapabilities(streaming=True),
-        skills=[],
-        default_input_modes=[],
-        default_output_modes=[],
-        preferred_transport='jsonrpc',
-    )
-
-
-@pytest.fixture
-def grpc_agent_card() -> AgentCard:
-    return AgentCard(
-        name='Test Agent',
-        description='A test agent',
-        url='http://test.agent.com/rpc',
-        version='1.0.0',
-        capabilities=AgentCapabilities(streaming=True),
-        skills=[],
-        default_input_modes=[],
-        default_output_modes=[],
-        preferred_transport='grpc',
-    )
-
-
-@pytest.mark.asyncio
-async def test_a2a_client_send_message(
-    mock_httpx_client: AsyncMock, jsonrpc_agent_card: AgentCard
-):
-    """Tests for the legacy client compatibility layer."""
-
-
 from unittest.mock import AsyncMock
 
 import pytest
 
-from a2a.types import (
-    AgentCard,
-)
+from a2a.types import AgentCard
 
 
 @pytest.fixture
@@ -164,7 +91,7 @@ async def test_a2a_grpc_client_get_task(
         status=TaskStatus(state=TaskState.working),
     )
 
-    client.get_task.return_value = mock_response_task
+    client.get_task = AsyncMock(return_value=mock_response_task)
 
     params = TaskQueryParams(id='task-456')
     response = await client.get_task(params)