Add tests for BaseClient and ClientTaskManager

Author: mikeas1

src/a2a/client/client_factory.py

@@ -23,7 +23,7 @@
 try:
     from a2a.client.transports.grpc import GrpcTransport
 except ImportError:
-    GrpcTransport = None  # type: ignore
+    GrpcTransport = None  # type: ignore # pyright: ignore
 
 
 logger = logging.getLogger(__name__)

tests/client/test_base_client.py

@@ -0,0 +1,118 @@
+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 import (
+    AgentCapabilities,
+    AgentCard,
+    Message,
+    Part,
+    Role,
+    Task,
+    TaskState,
+    TaskStatus,
+    TextPart,
+)
+
+
+@pytest.fixture
+def mock_transport():
+    transport = AsyncMock(spec=ClientTransport)
+    return transport
+
+
+@pytest.fixture
+def sample_agent_card():
+    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():
+    return Message(
+        role=Role.user,
+        message_id='msg-1',
+        parts=[Part(root=TextPart(text='Hello'))],
+    )
+
+
+@pytest.fixture
+def base_client(sample_agent_card, mock_transport):
+    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
+):
+    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()
+
+    events = [event async for event in base_client.send_message(sample_message)]
+
+    mock_transport.send_message_streaming.assert_called_once()
+    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
+):
+    base_client._config.streaming = False
+    mock_transport.send_message.return_value = Task(
+        id='task-456',
+        context_id='ctx-789',
+        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-456'
+
+
+@pytest.mark.asyncio
+async def test_send_message_non_streaming_agent_capability_false(
+    base_client: BaseClient, mock_transport: MagicMock, sample_message: Message
+):
+    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'

tests/client/test_client_task_manager.py

@@ -0,0 +1,163 @@
+import pytest
+from unittest.mock import AsyncMock, Mock, patch
+from a2a.client.client_task_manager import ClientTaskManager
+from a2a.client.errors import A2AClientInvalidArgsError, A2AClientInvalidStateError
+from a2a.types import (
+    Task,
+    TaskStatus,
+    TaskState,
+    TaskStatusUpdateEvent,
+    TaskArtifactUpdateEvent,
+    Message,
+    Role,
+    Part,
+    TextPart,
+    Artifact,
+)
+
+
+@pytest.fixture
+def task_manager():
+    return ClientTaskManager()
+
+
+@pytest.fixture
+def sample_task():
+    return Task(
+        id="task123",
+        context_id="context456",
+        status=TaskStatus(state=TaskState.working),
+        history=[],
+        artifacts=[],
+    )
+
+
+@pytest.fixture
+def sample_message():
+    return Message(
+        message_id="msg1",
+        role=Role.user,
+        parts=[Part(root=TextPart(text="Hello"))],
+    )
+
+
+def test_get_task_no_task_id_returns_none(task_manager: ClientTaskManager):
+    assert task_manager.get_task() is None
+
+
+def test_get_task_or_raise_no_task_raises_error(task_manager: ClientTaskManager):
+    with pytest.raises(A2AClientInvalidStateError, match="no current Task"):
+        task_manager.get_task_or_raise()
+
+
+@pytest.mark.asyncio
+async def test_save_task_event_with_task(
+    task_manager: ClientTaskManager, sample_task: Task
+):
+    await task_manager.save_task_event(sample_task)
+    assert task_manager.get_task() == sample_task
+    assert task_manager._task_id == sample_task.id
+    assert task_manager._context_id == sample_task.context_id
+
+
+@pytest.mark.asyncio
+async def test_save_task_event_with_task_already_set_raises_error(
+    task_manager: ClientTaskManager, sample_task: Task
+):
+    await task_manager.save_task_event(sample_task)
+    with pytest.raises(
+        A2AClientInvalidArgsError,
+        match="Task is already set, create new manager for new tasks.",
+    ):
+        await task_manager.save_task_event(sample_task)
+
+
+@pytest.mark.asyncio
+async def test_save_task_event_with_status_update(
+    task_manager: ClientTaskManager, sample_task: Task, sample_message: Message
+):
+    await task_manager.save_task_event(sample_task)
+    status_update = TaskStatusUpdateEvent(
+        task_id=sample_task.id,
+        context_id=sample_task.context_id,
+        status=TaskStatus(state=TaskState.completed, message=sample_message),
+        final=True,
+    )
+    updated_task = await task_manager.save_task_event(status_update)
+    assert updated_task.status.state == TaskState.completed
+    assert updated_task.history == [sample_message]
+
+
+@pytest.mark.asyncio
+async def test_save_task_event_with_artifact_update(
+    task_manager: ClientTaskManager, sample_task: Task
+):
+    await task_manager.save_task_event(sample_task)
+    artifact = Artifact(
+        artifact_id="art1", parts=[Part(root=TextPart(text="artifact content"))]
+    )
+    artifact_update = TaskArtifactUpdateEvent(
+        task_id=sample_task.id,
+        context_id=sample_task.context_id,
+        artifact=artifact,
+    )
+
+    with patch("a2a.client.client_task_manager.append_artifact_to_task") as mock_append:
+        updated_task = await task_manager.save_task_event(artifact_update)
+        mock_append.assert_called_once_with(updated_task, artifact_update)
+
+
+@pytest.mark.asyncio
+async def test_save_task_event_creates_task_if_not_exists(
+    task_manager: ClientTaskManager,
+):
+    status_update = TaskStatusUpdateEvent(
+        task_id="new_task",
+        context_id="new_context",
+        status=TaskStatus(state=TaskState.working),
+        final=False,
+    )
+    updated_task = await task_manager.save_task_event(status_update)
+    assert updated_task is not None
+    assert updated_task.id == "new_task"
+    assert updated_task.status.state == TaskState.working
+
+
+@pytest.mark.asyncio
+async def test_process_with_task_event(task_manager: ClientTaskManager, sample_task: Task):
+    with patch.object(
+        task_manager, "save_task_event", new_callable=AsyncMock
+    ) as mock_save:
+        await task_manager.process(sample_task)
+        mock_save.assert_called_once_with(sample_task)
+
+
+@pytest.mark.asyncio
+async def test_process_with_non_task_event(task_manager: ClientTaskManager):
+    with patch.object(
+        task_manager, "save_task_event", new_callable=Mock
+    ) as mock_save:
+        non_task_event = "not a task event"
+        await task_manager.process(non_task_event)
+        mock_save.assert_not_called()
+
+
+def test_update_with_message(
+    task_manager: ClientTaskManager, sample_task: Task, sample_message: Message
+):
+    updated_task = task_manager.update_with_message(sample_message, sample_task)
+    assert updated_task.history == [sample_message]
+
+
+def test_update_with_message_moves_status_message(
+    task_manager: ClientTaskManager, sample_task: Task, sample_message: Message
+):
+    status_message = Message(
+        message_id="status_msg",
+        role=Role.agent,
+        parts=[Part(root=TextPart(text="Status"))],
+    )
+    sample_task.status.message = status_message
+    updated_task = task_manager.update_with_message(sample_message, sample_task)
+    assert updated_task.history == [status_message, sample_message]
+    assert updated_task.status.message is None