test_base_client.py

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from a2a.client.base_client import BaseClient
from a2a.client.client import ClientConfig
from a2a.client.transports.base import ClientTransport
from a2a.types.a2a_pb2 import (
    AgentCapabilities,
    AgentCard,
    AgentInterface,
    CancelTaskRequest,
    TaskPushNotificationConfig,
    DeleteTaskPushNotificationConfigRequest,
    GetTaskPushNotificationConfigRequest,
    GetTaskRequest,
    ListTaskPushNotificationConfigsRequest,
    ListTaskPushNotificationConfigsResponse,
    ListTasksRequest,
    ListTasksResponse,
    Message,
    Part,
    Role,
    SendMessageConfiguration,
    SendMessageRequest,
    SendMessageResponse,
    StreamResponse,
    SubscribeToTaskRequest,
    Task,
    TaskPushNotificationConfig,
    TaskState,
    TaskStatus,
)


@pytest.fixture
def mock_transport() -> AsyncMock:
    return AsyncMock(spec=ClientTransport)


@pytest.fixture
def sample_agent_card() -> AgentCard:
    return AgentCard(
        name='Test Agent',
        description='An agent for testing',
        supported_interfaces=[
            AgentInterface(url='http://test.com', protocol_binding='HTTP+JSON')
        ],
        version='1.0',
        capabilities=AgentCapabilities(streaming=True),
        default_input_modes=['text/plain'],
        default_output_modes=['text/plain'],
        skills=[],
    )


@pytest.fixture
def sample_message() -> Message:
    return Message(
        role=Role.ROLE_USER,
        message_id='msg-1',
        parts=[Part(text='Hello')],
    )


@pytest.fixture
def base_client(
    sample_agent_card: AgentCard, mock_transport: AsyncMock
) -> BaseClient:
    config = ClientConfig(streaming=True)
    return BaseClient(
        card=sample_agent_card,
        config=config,
        transport=mock_transport,
        interceptors=[],
    )


class TestClientTransport:
    @pytest.mark.asyncio
    async def test_transport_async_context_manager(self) -> None:
        with (
            patch.object(ClientTransport, '__abstractmethods__', set()),
            patch.object(ClientTransport, 'close', new_callable=AsyncMock),
        ):
            transport = ClientTransport()
            async with transport as t:
                assert t is transport
                transport.close.assert_not_awaited()
            transport.close.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_transport_async_context_manager_on_exception(self) -> None:
        with (
            patch.object(ClientTransport, '__abstractmethods__', set()),
            patch.object(ClientTransport, 'close', new_callable=AsyncMock),
        ):
            transport = ClientTransport()
            with pytest.raises(RuntimeError, match='boom'):
                async with transport:
                    raise RuntimeError('boom')
            transport.close.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_base_client_async_context_manager(
        self, base_client: BaseClient, mock_transport: AsyncMock
    ) -> None:
        async with base_client as client:
            assert client is base_client
            mock_transport.close.assert_not_awaited()
        mock_transport.close.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_base_client_async_context_manager_on_exception(
        self, base_client: BaseClient, mock_transport: AsyncMock
    ) -> None:
        with pytest.raises(RuntimeError, match='boom'):
            async with base_client:
                raise RuntimeError('boom')
        mock_transport.close.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_send_message_streaming(
        self,
        base_client: BaseClient,
        mock_transport: MagicMock,
        sample_message: Message,
    ) -> None:
        async def create_stream(*args, **kwargs):
            task = Task(
                id='task-123',
                context_id='ctx-456',
                status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED),
            )
            stream_response = StreamResponse()
            stream_response.task.CopyFrom(task)
            yield stream_response

        mock_transport.send_message_streaming.return_value = create_stream()

        meta = {'test': 1}
        request = SendMessageRequest(message=sample_message, metadata=meta)
        stream = base_client.send_message(request)
        events = [event async for event in stream]

        mock_transport.send_message_streaming.assert_called_once()
        assert (
            mock_transport.send_message_streaming.call_args[0][0].metadata
            == meta
        )
        assert not mock_transport.send_message.called
        assert len(events) == 1
        response = events[0]
        assert response.task.id == 'task-123'

    @pytest.mark.asyncio
    async def test_send_message_non_streaming(
        self,
        base_client: BaseClient,
        mock_transport: MagicMock,
        sample_message: Message,
    ) -> None:
        base_client._config.streaming = False
        task = Task(
            id='task-456',
            context_id='ctx-789',
            status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED),
        )
        response = SendMessageResponse()
        response.task.CopyFrom(task)
        mock_transport.send_message.return_value = response

        meta = {'test': 1}
        request = SendMessageRequest(message=sample_message, metadata=meta)
        stream = base_client.send_message(request)
        events = [event async for event in stream]

        mock_transport.send_message.assert_called_once()
        assert mock_transport.send_message.call_args[0][0].metadata == meta
        assert not mock_transport.send_message_streaming.called
        assert len(events) == 1
        response = events[0]
        assert response.task.id == 'task-456'

    @pytest.mark.asyncio
    async def test_send_message_non_streaming_agent_capability_false(
        self,
        base_client: BaseClient,
        mock_transport: MagicMock,
        sample_message: Message,
    ) -> None:
        base_client._card.capabilities.streaming = False
        task = Task(
            id='task-789',
            context_id='ctx-101',
            status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED),
        )
        response = SendMessageResponse()
        response.task.CopyFrom(task)
        mock_transport.send_message.return_value = response

        request = SendMessageRequest(message=sample_message)
        events = [event async for event in base_client.send_message(request)]

        mock_transport.send_message.assert_called_once()
        assert not mock_transport.send_message_streaming.called
        assert len(events) == 1
        response = events[0]
        assert response.task.id == 'task-789'

    @pytest.mark.asyncio
    async def test_send_message_does_not_mutate_request(
        self,
        base_client: BaseClient,
        mock_transport: MagicMock,
        sample_message: Message,
    ):
        base_client._config.streaming = False
        base_client._config.polling = True
        base_client._config.accepted_output_modes = ['application/json']
        base_client._config.push_notification_configs = [
            TaskPushNotificationConfig(
                task_id='task-1',
            )
        ]

        task = Task(
            id='task-no-mutate',
            context_id='ctx-no-mutate',
            status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED),
        )
        response = SendMessageResponse()
        response.task.CopyFrom(task)
        mock_transport.send_message.return_value = response

        request = SendMessageRequest(message=sample_message)

        original = SendMessageRequest()
        original.CopyFrom(request)

        events = [event async for event in base_client.send_message(request)]
        assert len(events) == 1

        assert request == original

    @pytest.mark.asyncio
    async def test_send_message_callsite_config_overrides_non_streaming(
        self,
        base_client: BaseClient,
        mock_transport: MagicMock,
        sample_message: Message,
    ):
        base_client._config.streaming = False
        task = Task(
            id='task-cfg-ns-1',
            context_id='ctx-cfg-ns-1',
            status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED),
        )
        response = SendMessageResponse()
        response.task.CopyFrom(task)
        mock_transport.send_message.return_value = response

        cfg = SendMessageConfiguration(
            history_length=2,
            return_immediately=True,
            accepted_output_modes=['application/json'],
        )
        request = SendMessageRequest(message=sample_message, configuration=cfg)
        events = [event async for event in base_client.send_message(request)]

        mock_transport.send_message.assert_called_once()
        assert not mock_transport.send_message_streaming.called
        assert len(events) == 1
        response = events[0]
        assert response.task.id == 'task-cfg-ns-1'

        params = mock_transport.send_message.call_args[0][0]
        assert params.configuration.history_length == 2
        assert params.configuration.return_immediately is True
        assert params.configuration.accepted_output_modes == [
            'application/json'
        ]

    @pytest.mark.asyncio
    async def test_send_message_callsite_config_overrides_streaming(
        self,
        base_client: BaseClient,
        mock_transport: MagicMock,
        sample_message: Message,
    ):
        base_client._config.streaming = True
        base_client._card.capabilities.streaming = True

        async def create_stream(*args, **kwargs):
            task = Task(
                id='task-cfg-s-1',
                context_id='ctx-cfg-s-1',
                status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED),
            )
            stream_response = StreamResponse()
            stream_response.task.CopyFrom(task)
            yield stream_response

        mock_transport.send_message_streaming.return_value = create_stream()

        cfg = SendMessageConfiguration(
            history_length=0,
            accepted_output_modes=['text/plain'],
        )
        request = SendMessageRequest(message=sample_message, configuration=cfg)
        events = [event async for event in base_client.send_message(request)]

        mock_transport.send_message_streaming.assert_called_once()
        assert not mock_transport.send_message.called
        assert len(events) == 1
        response = events[0]
        assert response.task.id == 'task-cfg-s-1'

        params = mock_transport.send_message_streaming.call_args[0][0]
        assert params.configuration.history_length == 0
        assert params.configuration.return_immediately is False
        assert params.configuration.accepted_output_modes == ['text/plain']