Add unit tests for fastapi_app sendMessage api

Author: aneeshgarg

Gemini.md

@@ -0,0 +1,27 @@
+**A2A specification:** https://a2a-protocol.org/latest/specification/
+
+## Project frameworks
+- uv as package manager
+
+## How to run all tests
+1. If dependencies are not installed install them using following command
+   ```
+   uv sync --all-extras 
+   ```
+
+2. Run tests
+  ```
+  uv run pytest
+  ```
+
+## Other instructions
+1. Whenever writing python code, write types as well.
+2. After making the changes run ruff to check and fix the formatting issues
+   ```
+   uv run ruff check --fix
+   ```
+3. Run mypy type checkers to check for type errors
+   ```
+   uv run mypy
+   ```
+4. Run the unit tests to make sure that none of the unit tests are broken.

pyproject.toml

@@ -8,7 +8,7 @@ authors = [{ name = "Google LLC", email = "[email protected]" }]
 requires-python = ">=3.10"
 keywords = ["A2A", "A2A SDK", "A2A Protocol", "Agent2Agent", "Agent 2 Agent"]
 dependencies = [
-  "fastapi>=0.115.2",
+  "fastapi>=0.116.1",
   "httpx>=0.28.1",
   "httpx-sse>=0.4.0",
   "opentelemetry-api>=1.33.0",
@@ -93,6 +93,7 @@ dev = [
   "pyupgrade",
   "autoflake",
   "no_implicit_optional",
+  "trio",
 ]
 
 [[tool.uv.index]]

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

@@ -6,10 +6,11 @@
 from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
 from typing import Any
 
+from google.protobuf import message as message_pb2
 from pydantic import ValidationError
 from sse_starlette.sse import EventSourceResponse
 from starlette.requests import Request
-from starlette.responses import JSONResponse
+from starlette.responses import JSONResponse, Response
 
 from a2a.server.apps.jsonrpc import (
     CallContextBuilder,
@@ -111,9 +112,11 @@ def _handle_error(self, error: Exception) -> JSONResponse:
 
     async def _handle_request(
         self,
-        method: Callable[[Request, ServerCallContext], Awaitable[str]],
+        method: Callable[
+            [Request, ServerCallContext], Awaitable[dict[str, Any]]
+        ],
         request: Request,
-    ) -> JSONResponse:
+    ) -> Response:
         try:
             call_context = self._context_builder.build(request)
             response = await method(request, call_context)
@@ -123,7 +126,9 @@ async def _handle_request(
 
     async def _handle_streaming_request(
         self,
-        method: Callable[[Request, ServerCallContext], AsyncIterator[str]],
+        method: Callable[
+            [Request, ServerCallContext], AsyncIterator[message_pb2.Message]
+        ],
         request: Request,
     ) -> EventSourceResponse:
         try:

src/a2a/server/request_handlers/rest_handler.py

@@ -1,8 +1,9 @@
 import logging
 
 from collections.abc import AsyncIterable
+from typing import Any
 
-from google.protobuf.json_format import MessageToJson, Parse
+from google.protobuf.json_format import MessageToDict, MessageToJson, Parse
 from starlette.requests import Request
 
 from a2a.grpc import a2a_pb2
@@ -60,7 +61,7 @@ async def on_message_send(
         self,
         request: Request,
         context: ServerCallContext,
-    ) -> str:
+    ) -> dict[str, Any]:
         """Handles the 'message/send' REST method.
 
         Args:
@@ -85,7 +86,7 @@ async def on_message_send(
             task_or_message = await self.request_handler.on_message_send(
                 a2a_request, context
             )
-            return MessageToJson(
+            return MessageToDict(
                 proto_utils.ToProto.task_or_message(task_or_message)
             )
         except ServerError as e:

tests/server/apps/rest/test_fastapi_app.py

@@ -0,0 +1,148 @@
+from collections.abc import AsyncGenerator
+from unittest.mock import MagicMock
+import logging
+
+from google.protobuf import json_format
+import pytest
+from fastapi import FastAPI
+from httpx import ASGITransport, AsyncClient
+
+from a2a.grpc import a2a_pb2
+from a2a.server.apps.rest.fastapi_app import A2ARESTFastAPIApplication
+from a2a.server.request_handlers.request_handler import RequestHandler
+from a2a.types import (
+    AgentCard,
+    Message,
+    Part,
+    Role,
+    SendMessageRequest,
+    SendMessageResponse,
+    SendMessageSuccessResponse,
+    Task,
+    TaskState,
+    TaskStatus,
+    TextPart,
+)
+from a2a.utils import proto_utils
+
+
+logger = logging.getLogger(__name__)
+
+
+@pytest.fixture
+async def agent_card() -> AgentCard:
+    mock_agent_card = MagicMock(spec=AgentCard)
+    mock_agent_card.url = 'http://mockurl.com'
+    mock_agent_card.supports_authenticated_extended_card = False
+    return mock_agent_card
+
+
+@pytest.fixture
+async def request_handler() -> RequestHandler:
+    return MagicMock(spec=RequestHandler)
+
+
+@pytest.fixture
+async def app(
+    agent_card: AgentCard, request_handler: RequestHandler
+) -> FastAPI:
+    """Builds the FastAPI application for testing."""
+
+    return A2ARESTFastAPIApplication(agent_card, request_handler).build(
+        agent_card_url='/well-known/agent.json', rpc_url=''
+    )
+
+
+@pytest.fixture
+async def client(app: FastAPI) -> AsyncClient:
+    return AsyncClient(
+        transport=ASGITransport(app=app), base_url='http://testapp'
+    )
+
+
+@pytest.mark.anyio
+async def test_send_message_success_message(
+    client: AsyncClient, request_handler: MagicMock
+) -> None:
+    expected_response = a2a_pb2.SendMessageResponse(
+        msg=a2a_pb2.Message(
+            message_id='test',
+            role=a2a_pb2.Role.ROLE_AGENT,
+            content=[
+                a2a_pb2.Part(text='response message'),
+            ],
+        ),
+    )
+    request_handler.on_message_send.return_value = Message(
+        message_id='test',
+        role=Role.agent,
+        parts=[Part(TextPart(text='response message'))],
+    )
+
+    request = a2a_pb2.SendMessageRequest(
+        request=a2a_pb2.Message(),
+        configuration=a2a_pb2.SendMessageConfiguration(),
+    )
+    # To see log output, run pytest with '--log-cli=true --log-cli-level=INFO'
+    response = await client.post(
+        '/v1/message:send', json=json_format.MessageToDict(request)
+    )
+    # request should always be successful
+    response.raise_for_status()
+
+    actual_response = a2a_pb2.SendMessageResponse()
+    json_format.Parse(response.text, actual_response)
+    assert expected_response == actual_response
+
+
+@pytest.mark.anyio
+async def test_send_message_success_task(
+    client: AsyncClient, request_handler: MagicMock
+) -> None:
+    expected_response = a2a_pb2.SendMessageResponse(
+        task=a2a_pb2.Task(
+            id='test_task_id',
+            context_id='test_context_id',
+            status=a2a_pb2.TaskStatus(
+                state=a2a_pb2.TaskState.TASK_STATE_COMPLETED,
+                update=a2a_pb2.Message(
+                    message_id='test',
+                    role=a2a_pb2.ROLE_AGENT,
+                    content=[
+                        a2a_pb2.Part(text='response task message'),
+                    ],
+                ),
+            ),
+        ),
+    )
+    request_handler.on_message_send.return_value = Task(
+        id='test_task_id',
+        context_id='test_context_id',
+        status=TaskStatus(
+            state=TaskState.completed,
+            message=Message(
+                message_id='test',
+                role=Role.agent,
+                parts=[Part(TextPart(text='response task message'))],
+            ),
+        ),
+    )
+
+    request = a2a_pb2.SendMessageRequest(
+        request=a2a_pb2.Message(),
+        configuration=a2a_pb2.SendMessageConfiguration(),
+    )
+    # To see log output, run pytest with '--log-cli=true --log-cli-level=INFO'
+    response = await client.post(
+        '/v1/message:send', json=json_format.MessageToDict(request)
+    )
+    # request should always be successful
+    response.raise_for_status()
+
+    actual_response = a2a_pb2.SendMessageResponse()
+    json_format.Parse(response.text, actual_response)
+    assert expected_response == actual_response
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])