test: Add new tests to increase coverage (#195)

Co-authored-by: kthota-g <kcthota@google.com>

Author: holtskinner

.github/actions/spelling/excludes.txt

@@ -88,4 +88,5 @@
 CHANGELOG.md
 noxfile.py
 ^src/a2a/grpc/
+^tests/
 .pre-commit-config.yaml

.github/workflows/unit-tests.yml

@@ -28,6 +28,6 @@ jobs:
       - name: Install dependencies
         run: uv sync --dev
       - name: Run tests and check coverage
-        run: uv run pytest --cov=a2a --cov-report=xml --cov-fail-under=85
+        run: uv run pytest --cov=a2a --cov-report=xml --cov-fail-under=90
       - name: Show coverage summary in log
         run: uv run coverage report

src/a2a/utils/proto_utils.py

@@ -143,11 +143,16 @@ def authentication_info(
     def push_notification_config(
         cls, config: types.PushNotificationConfig
     ) -> a2a_pb2.PushNotificationConfig:
+        auth_info = (
+            ToProto.authentication_info(config.authentication)
+            if config.authentication
+            else None
+        )
         return a2a_pb2.PushNotificationConfig(
             id=config.id or '',
             url=config.url,
             token=config.token,
-            authentication=ToProto.authentication_info(config.authentication),
+            authentication=auth_info,
         )
 
     @classmethod
@@ -185,7 +190,9 @@ def message_send_configuration(
             accepted_output_modes=list(config.acceptedOutputModes),
             push_notification=ToProto.push_notification_config(
                 config.pushNotificationConfig
-            ),
+            )
+            if config.pushNotificationConfig
+            else None,
             history_length=config.historyLength,
             blocking=config.blocking or False,
         )
@@ -335,7 +342,7 @@ def security_scheme(
             return a2a_pb2.SecurityScheme(
                 api_key_security_scheme=a2a_pb2.APIKeySecurityScheme(
                     description=scheme.root.description,
-                    location=scheme.root.in_,
+                    location=scheme.root.in_.value,
                     name=scheme.root.name,
                 )
             )
@@ -548,7 +555,9 @@ def push_notification_config(
             id=config.id,
             url=config.url,
             token=config.token,
-            authentication=FromProto.authentication_info(config.authentication),
+            authentication=FromProto.authentication_info(config.authentication)
+            if config.HasField('authentication')
+            else None,
         )
 
     @classmethod
@@ -568,7 +577,9 @@ def message_send_configuration(
             acceptedOutputModes=list(config.accepted_output_modes),
             pushNotificationConfig=FromProto.push_notification_config(
                 config.push_notification
-            ),
+            )
+            if config.HasField('push_notification')
+            else None,
             historyLength=config.history_length,
             blocking=config.blocking,
         )
@@ -720,7 +731,7 @@ def security_scheme(
                 root=types.APIKeySecurityScheme(
                     description=scheme.api_key_security_scheme.description,
                     name=scheme.api_key_security_scheme.name,
-                    in_=scheme.api_key_security_scheme.location,  # type: ignore[call-arg]
+                    in_=types.In(scheme.api_key_security_scheme.location), # type: ignore[call-arg]
                 )
             )
         if scheme.HasField('http_auth_security_scheme'):

tests/client/test_grpc_client.py

@@ -0,0 +1,145 @@
+from unittest.mock import AsyncMock
+
+import pytest
+
+from a2a.client import A2AGrpcClient
+from a2a.grpc import a2a_pb2, a2a_pb2_grpc
+from a2a.types import (
+    AgentCapabilities,
+    AgentCard,
+    Message,
+    MessageSendParams,
+    Part,
+    Role,
+    Task,
+    TaskIdParams,
+    TaskQueryParams,
+    TaskState,
+    TaskStatus,
+    TextPart,
+)
+from a2a.utils import proto_utils
+
+
+# Fixtures
+@pytest.fixture
+def mock_grpc_stub() -> AsyncMock:
+    """Provides a mock gRPC stub with methods mocked."""
+    stub = AsyncMock(spec=a2a_pb2_grpc.A2AServiceStub)
+    stub.SendMessage = AsyncMock()
+    stub.SendStreamingMessage = AsyncMock()
+    stub.GetTask = AsyncMock()
+    stub.CancelTask = AsyncMock()
+    stub.CreateTaskPushNotification = AsyncMock()
+    stub.GetTaskPushNotification = AsyncMock()
+    return stub
+
+
+@pytest.fixture
+def sample_agent_card() -> AgentCard:
+    """Provides a minimal agent card for initialization."""
+    return AgentCard(
+        name='gRPC Test Agent',
+        description='Agent for testing gRPC client',
+        url='grpc://localhost:50051',
+        version='1.0',
+        capabilities=AgentCapabilities(streaming=True, pushNotifications=True),
+        defaultInputModes=['text/plain'],
+        defaultOutputModes=['text/plain'],
+        skills=[],
+    )
+
+
+@pytest.fixture
+def grpc_client(
+    mock_grpc_stub: AsyncMock, sample_agent_card: AgentCard
+) -> A2AGrpcClient:
+    """Provides an A2AGrpcClient instance."""
+    return A2AGrpcClient(grpc_stub=mock_grpc_stub, agent_card=sample_agent_card)
+
+
+@pytest.fixture
+def sample_message_send_params() -> MessageSendParams:
+    """Provides a sample MessageSendParams object."""
+    return MessageSendParams(
+        message=Message(
+            role=Role.user,
+            messageId='msg-1',
+            parts=[Part(root=TextPart(text='Hello'))],
+        )
+    )
+
+
+@pytest.fixture
+def sample_task() -> Task:
+    """Provides a sample Task object."""
+    return Task(
+        id='task-1',
+        contextId='ctx-1',
+        status=TaskStatus(state=TaskState.completed),
+    )
+
+
+@pytest.fixture
+def sample_message() -> Message:
+    """Provides a sample Message object."""
+    return Message(
+        role=Role.agent,
+        messageId='msg-response',
+        parts=[Part(root=TextPart(text='Hi there'))],
+    )
+
+
+@pytest.mark.asyncio
+async def test_send_message_task_response(
+    grpc_client: A2AGrpcClient,
+    mock_grpc_stub: AsyncMock,
+    sample_message_send_params: MessageSendParams,
+    sample_task: Task,
+):
+    """Test send_message that returns a Task."""
+    mock_grpc_stub.SendMessage.return_value = a2a_pb2.SendMessageResponse(
+        task=proto_utils.ToProto.task(sample_task)
+    )
+
+    response = await grpc_client.send_message(sample_message_send_params)
+
+    mock_grpc_stub.SendMessage.assert_awaited_once()
+    assert isinstance(response, Task)
+    assert response.id == sample_task.id
+
+
+@pytest.mark.asyncio
+async def test_get_task(
+    grpc_client: A2AGrpcClient, mock_grpc_stub: AsyncMock, sample_task: Task
+):
+    """Test retrieving a task."""
+    mock_grpc_stub.GetTask.return_value = proto_utils.ToProto.task(sample_task)
+    params = TaskQueryParams(id=sample_task.id)
+
+    response = await grpc_client.get_task(params)
+
+    mock_grpc_stub.GetTask.assert_awaited_once_with(
+        a2a_pb2.GetTaskRequest(name=f'tasks/{sample_task.id}')
+    )
+    assert response.id == sample_task.id
+
+
+@pytest.mark.asyncio
+async def test_cancel_task(
+    grpc_client: A2AGrpcClient, mock_grpc_stub: AsyncMock, sample_task: Task
+):
+    """Test cancelling a task."""
+    cancelled_task = sample_task.model_copy()
+    cancelled_task.status.state = TaskState.canceled
+    mock_grpc_stub.CancelTask.return_value = proto_utils.ToProto.task(
+        cancelled_task
+    )
+    params = TaskIdParams(id=sample_task.id)
+
+    response = await grpc_client.cancel_task(params)
+
+    mock_grpc_stub.CancelTask.assert_awaited_once_with(
+        a2a_pb2.CancelTaskRequest(name=f'tasks/{sample_task.id}')
+    )
+    assert response.status.state == TaskState.canceled

tests/server/agent_execution/test_context.py

@@ -10,6 +10,7 @@
     MessageSendParams,
     Task,
 )
+from a2a.utils.errors import ServerError
 
 
 class TestRequestContext:
@@ -165,6 +166,33 @@ def test_check_or_generate_context_id_with_existing_context_id(
         assert context.context_id == existing_id
         assert mock_params.message.contextId == existing_id
 
+    def test_init_raises_error_on_task_id_mismatch(
+        self, mock_params, mock_task
+    ):
+        """Test that an error is raised if provided task_id mismatches task.id."""
+        with pytest.raises(ServerError) as exc_info:
+            RequestContext(
+                request=mock_params, task_id='wrong-task-id', task=mock_task
+            )
+        assert 'bad task id' in str(exc_info.value.error.message)
+
+    def test_init_raises_error_on_context_id_mismatch(
+        self, mock_params, mock_task
+    ):
+        """Test that an error is raised if provided context_id mismatches task.contextId."""
+        # Set a valid task_id to avoid that error
+        mock_params.message.taskId = mock_task.id
+
+        with pytest.raises(ServerError) as exc_info:
+            RequestContext(
+                request=mock_params,
+                task_id=mock_task.id,
+                context_id='wrong-context-id',
+                task=mock_task,
+            )
+
+        assert 'bad context id' in str(exc_info.value.error.message)
+
     def test_with_related_tasks_provided(self, mock_task):
         """Test initialization with related tasks provided."""
         related_tasks = [mock_task, Mock(spec=Task)]

tests/server/apps/jsonrpc/test_jsonrpc_app.py

@@ -70,17 +70,20 @@ def test_jsonrpc_app_build_method_abstract_raises_typeerror(
         # Ensure 'supportsAuthenticatedExtendedCard' attribute exists
         mock_agent_card.supportsAuthenticatedExtendedCard = False
 
-        class AbstractTester(JSONRPCApplication):
-            # No 'build' method implemented
-            pass
-
-        # Instantiating an ABC subclass that doesn't implement all abstract methods raises TypeError
+        # This will fail at definition time if an abstract method is not implemented
         with pytest.raises(
             TypeError,
-            match="Can't instantiate abstract class AbstractTester with abstract method build",
+            match="Can't instantiate abstract class IncompleteJSONRPCApp with abstract method build",
         ):
-            # Using positional arguments for the abstract class constructor
-            AbstractTester(mock_handler, mock_agent_card)
+
+            class IncompleteJSONRPCApp(JSONRPCApplication):
+                # Intentionally not implementing 'build'
+                def some_other_method(self):
+                    pass
+
+            IncompleteJSONRPCApp(
+                agent_card=mock_agent_card, http_handler=mock_handler
+            )
 
 
 if __name__ == '__main__':

tests/server/request_handlers/test_grpc_handler.py

@@ -195,3 +195,87 @@ async def test_get_agent_card(
 
     assert response.name == sample_agent_card.name
     assert response.version == sample_agent_card.version
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'server_error, grpc_status_code, error_message_part',
+    [
+        (
+            ServerError(error=types.JSONParseError()),
+            grpc.StatusCode.INTERNAL,
+            'JSONParseError',
+        ),
+        (
+            ServerError(error=types.InvalidRequestError()),
+            grpc.StatusCode.INVALID_ARGUMENT,
+            'InvalidRequestError',
+        ),
+        (
+            ServerError(error=types.MethodNotFoundError()),
+            grpc.StatusCode.NOT_FOUND,
+            'MethodNotFoundError',
+        ),
+        (
+            ServerError(error=types.InvalidParamsError()),
+            grpc.StatusCode.INVALID_ARGUMENT,
+            'InvalidParamsError',
+        ),
+        (
+            ServerError(error=types.InternalError()),
+            grpc.StatusCode.INTERNAL,
+            'InternalError',
+        ),
+        (
+            ServerError(error=types.TaskNotFoundError()),
+            grpc.StatusCode.NOT_FOUND,
+            'TaskNotFoundError',
+        ),
+        (
+            ServerError(error=types.TaskNotCancelableError()),
+            grpc.StatusCode.UNIMPLEMENTED,
+            'TaskNotCancelableError',
+        ),
+        (
+            ServerError(error=types.PushNotificationNotSupportedError()),
+            grpc.StatusCode.UNIMPLEMENTED,
+            'PushNotificationNotSupportedError',
+        ),
+        (
+            ServerError(error=types.UnsupportedOperationError()),
+            grpc.StatusCode.UNIMPLEMENTED,
+            'UnsupportedOperationError',
+        ),
+        (
+            ServerError(error=types.ContentTypeNotSupportedError()),
+            grpc.StatusCode.UNIMPLEMENTED,
+            'ContentTypeNotSupportedError',
+        ),
+        (
+            ServerError(error=types.InvalidAgentResponseError()),
+            grpc.StatusCode.INTERNAL,
+            'InvalidAgentResponseError',
+        ),
+        (
+            ServerError(error=types.JSONRPCError(code=99, message='Unknown')),
+            grpc.StatusCode.UNKNOWN,
+            'Unknown error',
+        ),
+    ],
+)
+async def test_abort_context_error_mapping(
+    grpc_handler: GrpcHandler,
+    mock_request_handler: AsyncMock,
+    mock_grpc_context: AsyncMock,
+    server_error,
+    grpc_status_code,
+    error_message_part,
+):
+    mock_request_handler.on_get_task.side_effect = server_error
+    request_proto = a2a_pb2.GetTaskRequest(name='tasks/any')
+    await grpc_handler.GetTask(request_proto, mock_grpc_context)
+
+    mock_grpc_context.abort.assert_awaited_once()
+    call_args, _ = mock_grpc_context.abort.call_args
+    assert call_args[0] == grpc_status_code
+    assert error_message_part in call_args[1]