Increase code coverage for proto_utils.py

Author: holtskinner

tests/utils/test_proto_utils.py

@@ -1,6 +1,7 @@
 from unittest import mock
 
 import pytest
+from google.protobuf import struct_pb2
 
 from a2a import types
 from a2a.grpc import a2a_pb2
@@ -31,7 +32,30 @@ def sample_message() -> types.Message:
             ),
             types.Part(root=types.DataPart(data={'key': 'value'})),
         ],
-        metadata={'source': 'test'},
+        metadata={'source': 'test', 'nested': {'key': 'value'}},
+    )
+
+
+@pytest.fixture
+def sample_message_with_bytes() -> types.Message:
+    """A sample message that includes a file with raw bytes."""
+    return types.Message(
+        message_id='msg-2',
+        context_id='ctx-2',
+        task_id='task-2',
+        role=types.Role.user,
+        parts=[
+            types.Part(
+                root=types.FilePart(
+                    file=types.FileWithBytes(
+                        bytes='file content',
+                        name='bytes.txt',
+                        mime_type='text/plain',
+                    ),
+                )
+            )
+        ],
+        metadata={},
     )
 
 
@@ -47,6 +71,9 @@ def sample_task(sample_message: types.Message) -> types.Task:
         artifacts=[
             types.Artifact(
                 artifact_id='art-1',
+                name='artifact1',
+                description='An artifact',
+                metadata={'source': 'test'},
                 parts=[
                     types.Part(root=types.TextPart(text='Artifact content'))
                 ],
@@ -63,7 +90,16 @@ def sample_agent_card() -> types.AgentCard:
         url='http://localhost',
         version='1.0.0',
         capabilities=types.AgentCapabilities(
-            streaming=True, push_notifications=True
+            streaming=True,
+            push_notifications=True,
+            extensions=[
+                types.AgentExtension(
+                    uri='ext:uri',
+                    description='An extension',
+                    params={'key': 'value'},
+                    required=True,
+                )
+            ],
         ),
         default_input_modes=['text/plain'],
         default_output_modes=['text/plain'],
@@ -73,6 +109,9 @@ def sample_agent_card() -> types.AgentCard:
                 name='Test Skill',
                 description='A test skill',
                 tags=['test'],
+                examples=['example1'],
+                input_modes=['text/plain'],
+                output_modes=['text/markdown'],
             )
         ],
         provider=types.AgentProvider(
@@ -106,10 +145,48 @@ def sample_agent_card() -> types.AgentCard:
                     open_id_connect_url='http://oidc.url'
                 )
             ),
+            'mtls': types.SecurityScheme(
+                root=types.MutualTLSSecurityScheme(description='mTLS auth')
+            ),
         },
     )
 
 
+@pytest.fixture
+def sample_task_status_update() -> types.TaskStatusUpdateEvent:
+    """Sample TaskStatusUpdateEvent for testing."""
+    return types.TaskStatusUpdateEvent(
+        task_id='task-1',
+        context_id='ctx-1',
+        status=types.TaskStatus(
+            state=types.TaskState.completed,
+            message=types.Message(
+                message_id='msg-final',
+                role=types.Role.agent,
+                parts=[types.Part(root=types.TextPart(text='Done!'))],
+            ),
+        ),
+        metadata={'final_update': True},
+        final=True,
+    )
+
+
+@pytest.fixture
+def sample_task_artifact_update() -> types.TaskArtifactUpdateEvent:
+    """Sample TaskArtifactUpdateEvent for testing."""
+    return types.TaskArtifactUpdateEvent(
+        task_id='task-1',
+        context_id='ctx-1',
+        artifact=types.Artifact(
+            artifact_id='art-1',
+            parts=[types.Part(root=types.TextPart(text=' some data'))],
+        ),
+        metadata={'chunk': 2},
+        append=True,
+        last_chunk=False,
+    )
+
+
 # --- Test Cases ---
 
 
@@ -127,6 +204,11 @@ class FakePartType:
         with pytest.raises(ValueError, match='Unsupported part type'):
             proto_utils.ToProto.part(mock_part)
 
+    def test_update_event_unsupported_type(self):
+        """Test that ToProto.update_event raises ValueError for an unsupported type."""
+        with pytest.raises(ValueError, match='Unsupported event type'):
+            proto_utils.ToProto.update_event('not a valid event')
+
 
 class TestFromProto:
     def test_part_unsupported_type(self):
@@ -143,6 +225,14 @@ def test_task_query_params_invalid_name(self):
             proto_utils.FromProto.task_query_params(request)
         assert isinstance(exc_info.value.error, types.InvalidParamsError)
 
+    def test_task_push_notification_config_request_invalid_parent(self):
+        request = a2a_pb2.CreateTaskPushNotificationConfigRequest(
+            parent='invalid-parent-format'
+        )
+        with pytest.raises(ServerError) as exc_info:
+            proto_utils.FromProto.task_push_notification_config_request(request)
+        assert isinstance(exc_info.value.error, types.InvalidParamsError)
+
 
 class TestProtoUtils:
     def test_roundtrip_message(self, sample_message: types.Message):
@@ -158,6 +248,92 @@ def test_roundtrip_message(self, sample_message: types.Message):
         roundtrip_msg = proto_utils.FromProto.message(proto_msg)
         assert roundtrip_msg == sample_message
 
+    def test_roundtrip_message_with_bytes(
+        self, sample_message_with_bytes: types.Message
+    ):
+        """Test conversion of a Message with file bytes to proto and back."""
+        proto_msg = proto_utils.ToProto.message(sample_message_with_bytes)
+        assert isinstance(proto_msg, a2a_pb2.Message)
+        assert proto_msg.content[0].file.file_with_bytes == b'file content'
+
+        roundtrip_msg = proto_utils.FromProto.message(proto_msg)
+        assert roundtrip_msg == sample_message_with_bytes
+        # Also check metadata with empty dict
+        assert roundtrip_msg.metadata == {}
+
+    def test_roundtrip_task(self, sample_task: types.Task):
+        """Test conversion of Task to proto and back."""
+        proto_task = proto_utils.ToProto.task(sample_task)
+        assert isinstance(proto_task, a2a_pb2.Task)
+
+        roundtrip_task = proto_utils.FromProto.task(proto_task)
+        assert roundtrip_task == sample_task
+
+    def test_roundtrip_agent_card(self, sample_agent_card: types.AgentCard):
+        """Test conversion of AgentCard to proto and back."""
+        proto_card = proto_utils.ToProto.agent_card(sample_agent_card)
+        assert isinstance(proto_card, a2a_pb2.AgentCard)
+        assert proto_card.security_schemes['mtls'].HasField(
+            'mtls_security_scheme'
+        )
+
+        roundtrip_card = proto_utils.FromProto.agent_card(proto_card)
+        assert roundtrip_card == sample_agent_card
+
+    def test_roundtrip_stream_responses(
+        self,
+        sample_task_status_update: types.TaskStatusUpdateEvent,
+        sample_task_artifact_update: types.TaskArtifactUpdateEvent,
+        sample_task: types.Task,
+        sample_message: types.Message,
+    ):
+        """Test roundtrip conversion for all streamable event types."""
+        # TaskStatusUpdateEvent
+        proto_status_update = proto_utils.ToProto.stream_response(
+            sample_task_status_update
+        )
+        roundtrip_status_update = proto_utils.FromProto.stream_response(
+            proto_status_update
+        )
+        assert roundtrip_status_update == sample_task_status_update
+
+        # TaskArtifactUpdateEvent
+        proto_artifact_update = proto_utils.ToProto.stream_response(
+            sample_task_artifact_update
+        )
+        roundtrip_artifact_update = proto_utils.FromProto.stream_response(
+            proto_artifact_update
+        )
+        assert roundtrip_artifact_update == sample_task_artifact_update
+
+        # Task
+        proto_task = proto_utils.ToProto.stream_response(sample_task)
+        roundtrip_task = proto_utils.FromProto.stream_response(proto_task)
+        assert roundtrip_task == sample_task
+
+        # Message
+        proto_message = proto_utils.ToProto.stream_response(sample_message)
+        roundtrip_message = proto_utils.FromProto.stream_response(proto_message)
+        assert roundtrip_message == sample_message
+
+    def test_task_or_message_conversion(
+        self, sample_task: types.Task, sample_message: types.Message
+    ):
+        """Test ToProto and FromProto for task_or_message methods."""
+        # Message case
+        proto_msg_resp = proto_utils.ToProto.task_or_message(sample_message)
+        assert proto_msg_resp.HasField('msg')
+        assert not proto_msg_resp.HasField('task')
+        roundtrip_msg = proto_utils.FromProto.task_or_message(proto_msg_resp)
+        assert roundtrip_msg == sample_message
+
+        # Task case
+        proto_task_resp = proto_utils.ToProto.task_or_message(sample_task)
+        assert not proto_task_resp.HasField('msg')
+        assert proto_task_resp.HasField('task')
+        roundtrip_task = proto_utils.FromProto.task_or_message(proto_task_resp)
+        assert roundtrip_task == sample_task
+
     def test_enum_conversions(self):
         """Test conversions for all enum types."""
         assert (
@@ -168,6 +344,11 @@ def test_enum_conversions(self):
             proto_utils.FromProto.role(a2a_pb2.Role.ROLE_USER)
             == types.Role.user
         )
+        # Test unspecified role defaults to agent
+        assert (
+            proto_utils.FromProto.role(a2a_pb2.Role.ROLE_UNSPECIFIED)
+            == types.Role.agent
+        )
 
         for state in types.TaskState:
             if state not in (types.TaskState.unknown, types.TaskState.rejected):
@@ -232,18 +413,68 @@ def test_oauth_flows_conversion(self):
         )
         assert roundtrip_implicit.implicit is not None
 
+        roundtrip_auth_code = proto_utils.FromProto.oauth2_flows(
+            proto_auth_code_flow
+        )
+        assert roundtrip_auth_code.authorization_code is not None
+
+    def test_task_id_params_from_proto(self):
+        """Test successful parsing of task ID from request names."""
+        request = a2a_pb2.CancelTaskRequest(name='tasks/task-123')
+        params = proto_utils.FromProto.task_id_params(request)
+        assert params.id == 'task-123'
+
+        push_request = a2a_pb2.GetTaskPushNotificationConfigRequest(
+            name='tasks/task-456/pushNotificationConfigs/config-abc'
+        )
+        params = proto_utils.FromProto.task_id_params(push_request)
+        assert params.id == 'task-456'
+
     def test_task_id_params_from_proto_invalid_name(self):
         request = a2a_pb2.CancelTaskRequest(name='invalid-name-format')
         with pytest.raises(ServerError) as exc_info:
             proto_utils.FromProto.task_id_params(request)
         assert isinstance(exc_info.value.error, types.InvalidParamsError)
 
+    def test_task_push_config_from_proto(self):
+        """Test successful parsing of task push notification config."""
+        name = 'tasks/task-789/pushNotificationConfigs/config-xyz'
+        proto_config = a2a_pb2.TaskPushNotificationConfig(
+            name=name,
+            push_notification_config=a2a_pb2.PushNotificationConfig(id='cfg-1'),
+        )
+        config = proto_utils.FromProto.task_push_notification_config(
+            proto_config
+        )
+        assert config.task_id == 'task-789'
+        assert config.push_notification_config.id == 'cfg-1'
+
     def test_task_push_config_from_proto_invalid_parent(self):
         request = a2a_pb2.TaskPushNotificationConfig(name='invalid-name-format')
         with pytest.raises(ServerError) as exc_info:
             proto_utils.FromProto.task_push_notification_config(request)
         assert isinstance(exc_info.value.error, types.InvalidParamsError)
 
+    def test_from_proto_task_query_params(self):
+        """Test successful parsing of task query parameters."""
+        request = a2a_pb2.GetTaskRequest(
+            name='tasks/task-abc', history_length=50
+        )
+        params = proto_utils.FromProto.task_query_params(request)
+        assert params.id == 'task-abc'
+        assert params.history_length == 50
+
+    def test_dict_to_struct(self):
+        """Test the dict_to_struct utility function."""
+        py_dict = {'a': 1, 'b': 'hello', 'c': {'d': True}}
+        struct = proto_utils.dict_to_struct(py_dict)
+
+        assert isinstance(struct, struct_pb2.Struct)
+        assert struct['a'] == 1
+        assert struct['b'] == 'hello'
+        assert isinstance(struct['c'], struct_pb2.Struct)
+        assert struct['c']['d'] is True
+
     def test_none_handling(self):
         """Test that None inputs are handled gracefully."""
         assert proto_utils.ToProto.message(None) is None