test_proto_utils.py

from unittest import mock

import pytest

from a2a import types
from a2a.grpc import a2a_pb2
from a2a.utils import proto_utils
from a2a.utils.errors import ServerError


# --- Test Data ---


@pytest.fixture
def sample_message() -> types.Message:
    return types.Message(
        message_id='msg-1',
        context_id='ctx-1',
        task_id='task-1',
        role=types.Role.user,
        parts=[
            types.Part(root=types.TextPart(text='Hello')),
            types.Part(
                root=types.FilePart(
                    file=types.FileWithUri(
                        uri='file:///test.txt',
                        name='test.txt',
                        mime_type='text/plain',
                    ),
                )
            ),
            types.Part(root=types.DataPart(data={'key': 'value'})),
        ],
        metadata={'source': 'test'},
    )


@pytest.fixture
def sample_task(sample_message: types.Message) -> types.Task:
    return types.Task(
        id='task-1',
        context_id='ctx-1',
        status=types.TaskStatus(
            state=types.TaskState.working, message=sample_message
        ),
        history=[sample_message],
        artifacts=[
            types.Artifact(
                artifact_id='art-1',
                parts=[
                    types.Part(root=types.TextPart(text='Artifact content'))
                ],
            )
        ],
    )


@pytest.fixture
def sample_agent_card() -> types.AgentCard:
    return types.AgentCard(
        name='Test Agent',
        description='A test agent',
        url='http://localhost',
        version='1.0.0',
        capabilities=types.AgentCapabilities(
            streaming=True, push_notifications=True
        ),
        default_input_modes=['text/plain'],
        default_output_modes=['text/plain'],
        skills=[
            types.AgentSkill(
                id='skill1',
                name='Test Skill',
                description='A test skill',
                tags=['test'],
            )
        ],
        provider=types.AgentProvider(
            organization='Test Org', url='http://test.org'
        ),
        security=[{'oauth_scheme': ['read', 'write']}],
        security_schemes={
            'oauth_scheme': types.SecurityScheme(
                root=types.OAuth2SecurityScheme(
                    flows=types.OAuthFlows(
                        client_credentials=types.ClientCredentialsOAuthFlow(
                            token_url='http://token.url',
                            scopes={
                                'read': 'Read access',
                                'write': 'Write access',
                            },
                        )
                    )
                )
            ),
            'apiKey': types.SecurityScheme(
                root=types.APIKeySecurityScheme(
                    name='X-API-KEY', in_=types.In.header
                )
            ),
            'httpAuth': types.SecurityScheme(
                root=types.HTTPAuthSecurityScheme(scheme='bearer')
            ),
            'oidc': types.SecurityScheme(
                root=types.OpenIdConnectSecurityScheme(
                    open_id_connect_url='http://oidc.url'
                )
            ),
        },
    )


# --- Test Cases ---


class TestToProto:
    def test_part_unsupported_type(self):
        """Test that ToProto.part raises ValueError for an unsupported Part type."""

        class FakePartType:
            kind = 'fake'

        # Create a mock Part object that has a .root attribute pointing to the fake type
        mock_part = mock.MagicMock(spec=types.Part)
        mock_part.root = FakePartType()

        with pytest.raises(ValueError, match='Unsupported part type'):
            proto_utils.ToProto.part(mock_part)


class TestFromProto:
    def test_part_unsupported_type(self):
        """Test that FromProto.part raises ValueError for an unsupported part type in proto."""
        unsupported_proto_part = (
            a2a_pb2.Part()
        )  # An empty part with no oneof field set
        with pytest.raises(ValueError, match='Unsupported part type'):
            proto_utils.FromProto.part(unsupported_proto_part)

    def test_task_query_params_invalid_name(self):
        request = a2a_pb2.GetTaskRequest(name='invalid-name-format')
        with pytest.raises(ServerError) as exc_info:
            proto_utils.FromProto.task_query_params(request)
        assert isinstance(exc_info.value.error, types.InvalidParamsError)


class TestProtoUtils:
    def test_roundtrip_message(self, sample_message: types.Message):
        """Test conversion of Message to proto and back."""
        proto_msg = proto_utils.ToProto.message(sample_message)
        assert isinstance(proto_msg, a2a_pb2.Message)

        # Test file part handling
        assert proto_msg.content[1].file.file_with_uri == 'file:///test.txt'
        assert proto_msg.content[1].file.mime_type == 'text/plain'
        assert proto_msg.content[1].file.name == 'test.txt'

        roundtrip_msg = proto_utils.FromProto.message(proto_msg)
        assert roundtrip_msg == sample_message

    def test_enum_conversions(self):
        """Test conversions for all enum types."""
        assert (
            proto_utils.ToProto.role(types.Role.agent)
            == a2a_pb2.Role.ROLE_AGENT
        )
        assert (
            proto_utils.FromProto.role(a2a_pb2.Role.ROLE_USER)
            == types.Role.user
        )

        for state in types.TaskState:
            if state not in (
                types.TaskState.unknown,
                types.TaskState.rejected,
                types.TaskState.auth_required,
            ):
                proto_state = proto_utils.ToProto.task_state(state)
                assert proto_utils.FromProto.task_state(proto_state) == state

        # Test unknown state case
        assert (
            proto_utils.FromProto.task_state(
                a2a_pb2.TaskState.TASK_STATE_UNSPECIFIED
            )
            == types.TaskState.unknown
        )
        assert (
            proto_utils.ToProto.task_state(types.TaskState.unknown)
            == a2a_pb2.TaskState.TASK_STATE_UNSPECIFIED
        )

    def test_oauth_flows_conversion(self):
        """Test conversion of different OAuth2 flows."""
        # Test password flow
        password_flow = types.OAuthFlows(
            password=types.PasswordOAuthFlow(
                token_url='http://token.url', scopes={'read': 'Read'}
            )
        )
        proto_password_flow = proto_utils.ToProto.oauth2_flows(password_flow)
        assert proto_password_flow.HasField('password')

        # Test implicit flow
        implicit_flow = types.OAuthFlows(
            implicit=types.ImplicitOAuthFlow(
                authorization_url='http://auth.url', scopes={'read': 'Read'}
            )
        )
        proto_implicit_flow = proto_utils.ToProto.oauth2_flows(implicit_flow)
        assert proto_implicit_flow.HasField('implicit')

        # Test authorization code flow
        auth_code_flow = types.OAuthFlows(
            authorization_code=types.AuthorizationCodeOAuthFlow(
                authorization_url='http://auth.url',
                token_url='http://token.url',
                scopes={'read': 'read'},
            )
        )
        proto_auth_code_flow = proto_utils.ToProto.oauth2_flows(auth_code_flow)
        assert proto_auth_code_flow.HasField('authorization_code')

        # Test invalid flow
        with pytest.raises(ValueError):
            proto_utils.ToProto.oauth2_flows(types.OAuthFlows())

        # Test FromProto
        roundtrip_password = proto_utils.FromProto.oauth2_flows(
            proto_password_flow
        )
        assert roundtrip_password.password is not None

        roundtrip_implicit = proto_utils.FromProto.oauth2_flows(
            proto_implicit_flow
        )
        assert roundtrip_implicit.implicit is not None

    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_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_none_handling(self):
        """Test that None inputs are handled gracefully."""
        assert proto_utils.ToProto.message(None) is None
        assert proto_utils.ToProto.metadata(None) is None
        assert proto_utils.ToProto.provider(None) is None
        assert proto_utils.ToProto.security(None) is None
        assert proto_utils.ToProto.security_schemes(None) is None

    def test_metadata_conversion(self):
        """Test metadata conversion with various data types."""
        metadata = {
            'null_value': None,
            'bool_value': True,
            'int_value': 42,
            'float_value': 3.14,
            'string_value': 'hello',
            'dict_value': {'nested': 'dict', 'count': 10},
            'list_value': [1, 'two', 3.0, True, None],
            'tuple_value': (1, 2, 3),
            'complex_list': [
                {'name': 'item1', 'values': [1, 2, 3]},
                {'name': 'item2', 'values': [4, 5, 6]},
            ],
        }

        # Convert to proto
        proto_metadata = proto_utils.ToProto.metadata(metadata)
        assert proto_metadata is not None

        # Convert back to Python
        roundtrip_metadata = proto_utils.FromProto.metadata(proto_metadata)

        # Verify all values are preserved correctly
        assert roundtrip_metadata['null_value'] is None
        assert roundtrip_metadata['bool_value'] is True
        assert roundtrip_metadata['int_value'] == 42
        assert roundtrip_metadata['float_value'] == 3.14
        assert roundtrip_metadata['string_value'] == 'hello'
        assert roundtrip_metadata['dict_value']['nested'] == 'dict'
        assert roundtrip_metadata['dict_value']['count'] == 10
        assert roundtrip_metadata['list_value'] == [1, 'two', 3.0, True, None]
        assert roundtrip_metadata['tuple_value'] == [
            1,
            2,
            3,
        ]  # tuples become lists
        assert len(roundtrip_metadata['complex_list']) == 2
        assert roundtrip_metadata['complex_list'][0]['name'] == 'item1'

    def test_metadata_with_custom_objects(self):
        """Test metadata conversion with custom objects that need str() fallback."""

        class CustomObject:
            def __str__(self):
                return 'custom_object_str'

            def __repr__(self):
                return 'CustomObject()'

        metadata = {
            'custom_obj': CustomObject(),
            'list_with_custom': [1, CustomObject(), 'text'],
            'nested_custom': {'obj': CustomObject(), 'normal': 'value'},
        }

        # Convert to proto
        proto_metadata = proto_utils.ToProto.metadata(metadata)
        assert proto_metadata is not None

        # Convert back to Python
        roundtrip_metadata = proto_utils.FromProto.metadata(proto_metadata)

        # Custom objects should be converted to strings
        assert roundtrip_metadata['custom_obj'] == 'custom_object_str'
        assert roundtrip_metadata['list_with_custom'] == [
            1,
            'custom_object_str',
            'text',
        ]
        assert roundtrip_metadata['nested_custom']['obj'] == 'custom_object_str'
        assert roundtrip_metadata['nested_custom']['normal'] == 'value'

    def test_metadata_edge_cases(self):
        """Test metadata conversion with edge cases."""
        metadata = {
            'empty_dict': {},
            'empty_list': [],
            'zero': 0,
            'false': False,
            'empty_string': '',
            'unicode_string': 'string test',
            'large_number': 9999999999999999,
            'negative_number': -42,
            'float_precision': 0.123456789,
        }

        # Convert to proto and back
        proto_metadata = proto_utils.ToProto.metadata(metadata)
        roundtrip_metadata = proto_utils.FromProto.metadata(proto_metadata)

        # Verify edge cases are handled correctly
        assert roundtrip_metadata['empty_dict'] == {}
        assert roundtrip_metadata['empty_list'] == []
        assert roundtrip_metadata['zero'] == 0
        assert roundtrip_metadata['false'] is False
        assert roundtrip_metadata['empty_string'] == ''
        assert roundtrip_metadata['unicode_string'] == 'string test'
        assert roundtrip_metadata['large_number'] == 9999999999999999
        assert roundtrip_metadata['negative_number'] == -42
        assert abs(roundtrip_metadata['float_precision'] - 0.123456789) < 1e-10