test: improve coverage

Signed-off-by: Shingo OKAWA <[email protected]>

Author: ognis1205

src/a2a/client/grpc_client.py

@@ -64,7 +64,7 @@ async def send_message(
                 metadata=proto_utils.ToProto.metadata(request.metadata),
             )
         )
-        if response.task:
+        if response.HasField('task'):
             return proto_utils.FromProto.task(response.task)
         return proto_utils.FromProto.message(response.msg)
 
@@ -87,7 +87,7 @@ async def send_message_streaming(
             `TaskArtifactUpdateEvent` objects as they are received in the
             stream.
         """
-        stream = self.stub.SendStreamingMessage(
+        stream = await self.stub.SendStreamingMessage(
             a2a_pb2.SendMessageRequest(
                 request=proto_utils.ToProto.message(request.message),
                 configuration=proto_utils.ToProto.message_send_configuration(

tests/client/test_grpc_client.py

@@ -1,21 +1,30 @@
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, MagicMock
 
+import grpc
 import pytest
 
+from google.protobuf import struct_pb2
+
 from a2a.client import A2AGrpcClient
 from a2a.grpc import a2a_pb2, a2a_pb2_grpc
 from a2a.types import (
     AgentCapabilities,
     AgentCard,
+    Artifact,
     Message,
     MessageSendParams,
     Part,
+    PushNotificationAuthenticationInfo,
+    PushNotificationConfig,
     Role,
     Task,
+    TaskArtifactUpdateEvent,
     TaskIdParams,
+    TaskPushNotificationConfig,
     TaskQueryParams,
     TaskState,
     TaskStatus,
+    TaskStatusUpdateEvent,
     TextPart,
 )
 from a2a.utils import proto_utils
@@ -90,6 +99,78 @@ def sample_message() -> Message:
     )
 
 
+@pytest.fixture
+def sample_artifact() -> Artifact:
+    """Provides a sample Artifact object."""
+    return Artifact(
+        artifactId='artifact-1',
+        name='example.txt',
+        description='An example artifact',
+        parts=[Part(root=TextPart(text='Hi there'))],
+        metadata=struct_pb2.Struct(),
+        extensions=[],
+    )
+
+
+@pytest.fixture
+def sample_task_status_update_event() -> TaskStatusUpdateEvent:
+    """Provides a sample TaskStatusUpdateEvent."""
+    return TaskStatusUpdateEvent(
+        taskId='task-1',
+        contextId='ctx-1',
+        status=TaskStatus(state=TaskState.working),
+        final=False,
+        metadata=struct_pb2.Struct(),
+    )
+
+
+@pytest.fixture
+def sample_task_artifact_update_event(
+    sample_artifact,
+) -> TaskArtifactUpdateEvent:
+    """Provides a sample TaskArtifactUpdateEvent."""
+    return TaskArtifactUpdateEvent(
+        taskId='task-1',
+        contextId='ctx-1',
+        artifact=sample_artifact,
+        append=True,
+        last_chunk=True,
+        metadata=struct_pb2.Struct(),
+    )
+
+
+@pytest.fixture
+def sample_authentication_info() -> PushNotificationAuthenticationInfo:
+    """Provides a sample AuthenticationInfo object."""
+    return PushNotificationAuthenticationInfo(
+        schemes=['apikey', 'oauth2'], credentials='secret-token'
+    )
+
+
+@pytest.fixture
+def sample_push_notification_config(
+    sample_authentication_info: PushNotificationAuthenticationInfo,
+) -> PushNotificationConfig:
+    """Provides a sample PushNotificationConfig object."""
+    return PushNotificationConfig(
+        id='config-1',
+        url='https://example.com/notify',
+        token='example-token',
+        authentication=sample_authentication_info,
+    )
+
+
+@pytest.fixture
+def sample_task_push_notification_config(
+    sample_push_notification_config: PushNotificationConfig,
+) -> TaskPushNotificationConfig:
+    """Provides a sample TaskPushNotificationConfig object."""
+    return TaskPushNotificationConfig(
+        name='task-push-notification-config-name-1',
+        push_notification_config=sample_push_notification_config,
+    )
+
+
 @pytest.mark.asyncio
 async def test_send_message_task_response(
     grpc_client: A2AGrpcClient,
@@ -109,6 +190,76 @@ async def test_send_message_task_response(
     assert response.id == sample_task.id
 
 
+@pytest.mark.asyncio
+async def test_send_message_message_response(
+    grpc_client: A2AGrpcClient,
+    mock_grpc_stub: AsyncMock,
+    sample_message_send_params: MessageSendParams,
+    sample_message: Message,
+):
+    """Test send_message that returns a Message."""
+    mock_grpc_stub.SendMessage.return_value = a2a_pb2.SendMessageResponse(
+        msg=proto_utils.ToProto.message(sample_message)
+    )
+
+    response = await grpc_client.send_message(sample_message_send_params)
+
+    mock_grpc_stub.SendMessage.assert_awaited_once()
+    assert isinstance(response, Message)
+    assert response.messageId == sample_message.messageId
+
+
+@pytest.mark.asyncio
+async def test_send_message_streaming(
+    grpc_client: A2AGrpcClient,
+    mock_grpc_stub: AsyncMock,
+    sample_message_send_params: MessageSendParams,
+    sample_message: Message,
+    sample_task: Task,
+    sample_task_status_update_event: TaskStatusUpdateEvent,
+    sample_task_artifact_update_event: TaskArtifactUpdateEvent,
+):
+    """Test send_message_streaming that yields responses."""
+    stream = MagicMock()
+    stream.read = AsyncMock(
+        side_effect=[
+            a2a_pb2.StreamResponse(
+                msg=proto_utils.ToProto.message(sample_message)
+            ),
+            a2a_pb2.StreamResponse(task=proto_utils.ToProto.task(sample_task)),
+            a2a_pb2.StreamResponse(
+                status_update=proto_utils.ToProto.task_status_update_event(
+                    sample_task_status_update_event
+                )
+            ),
+            a2a_pb2.StreamResponse(
+                artifact_update=proto_utils.ToProto.task_artifact_update_event(
+                    sample_task_artifact_update_event
+                )
+            ),
+            grpc.aio.EOF,
+        ]
+    )
+    mock_grpc_stub.SendStreamingMessage.return_value = stream
+
+    responses = [
+        response
+        async for response in grpc_client.send_message_streaming(
+            sample_message_send_params
+        )
+    ]
+
+    mock_grpc_stub.SendStreamingMessage.assert_awaited_once()
+    assert isinstance(responses[0], Message)
+    assert responses[0].messageId == sample_message.messageId
+    assert isinstance(responses[1], Task)
+    assert responses[1].id == sample_task.id
+    assert isinstance(responses[2], TaskStatusUpdateEvent)
+    assert responses[2].taskId == sample_task_status_update_event.taskId
+    assert isinstance(responses[3], TaskArtifactUpdateEvent)
+    assert responses[3].taskId == sample_task_artifact_update_event.taskId
+
+
 @pytest.mark.asyncio
 async def test_get_task(
     grpc_client: A2AGrpcClient, mock_grpc_stub: AsyncMock, sample_task: Task