a2a-python/tests/client/test_base_client.py at 45ff87157a85215e904e319b315069d15e232d1c · a2aproject/a2a-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
from unittest.mock import AsyncMock, MagicMock

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 import (
    AgentCapabilities,
    AgentCard,
    Message,
    Part,
    Role,
    Task,
    TaskState,
    TaskStatus,
    TextPart,
)


@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',
        url='http://test.com',
        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.user,
        message_id='msg-1',
        parts=[Part(root=TextPart(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,
        consumers=[],
        middleware=[],
    )


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

    mock_transport.send_message_streaming.return_value = create_stream()

    meta = {'test': 1}
    stream = base_client.send_message(sample_message, request_metadata=meta)
    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
    assert events[0][0].id == 'task-123'


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

    meta = {'test': 1}
    stream = base_client.send_message(sample_message, request_metadata=meta)
    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
    assert events[0][0].id == 'task-456'


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

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

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