Formatting

Author: mikeas1

tests/client/test_client_task_manager.py

@@ -1,7 +1,10 @@
 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.client.errors import (
+    A2AClientInvalidArgsError,
+    A2AClientInvalidStateError,
+)
 from a2a.types import (
     Task,
     TaskStatus,
@@ -24,8 +27,8 @@ def task_manager():
 @pytest.fixture
 def sample_task():
     return Task(
-        id="task123",
-        context_id="context456",
+        id='task123',
+        context_id='context456',
         status=TaskStatus(state=TaskState.working),
         history=[],
         artifacts=[],
@@ -35,18 +38,20 @@ def sample_task():
 @pytest.fixture
 def sample_message():
     return Message(
-        message_id="msg1",
+        message_id='msg1',
         role=Role.user,
-        parts=[Part(root=TextPart(text="Hello"))],
+        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"):
+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()
 
 
@@ -67,7 +72,7 @@ async def test_save_task_event_with_task_already_set_raises_error(
     await task_manager.save_task_event(sample_task)
     with pytest.raises(
         A2AClientInvalidArgsError,
-        match="Task is already set, create new manager for new tasks.",
+        match='Task is already set, create new manager for new tasks.',
     ):
         await task_manager.save_task_event(sample_task)
 
@@ -94,15 +99,17 @@ async def test_save_task_event_with_artifact_update(
 ):
     await task_manager.save_task_event(sample_task)
     artifact = Artifact(
-        artifact_id="art1", parts=[Part(root=TextPart(text="artifact content"))]
+        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:
+    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)
 
@@ -112,21 +119,23 @@ 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",
+        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.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):
+async def test_process_with_task_event(
+    task_manager: ClientTaskManager, sample_task: Task
+):
     with patch.object(
-        task_manager, "save_task_event", new_callable=AsyncMock
+        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)
@@ -135,9 +144,9 @@ async def test_process_with_task_event(task_manager: ClientTaskManager, sample_t
 @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
+        task_manager, 'save_task_event', new_callable=Mock
     ) as mock_save:
-        non_task_event = "not a task event"
+        non_task_event = 'not a task event'
         await task_manager.process(non_task_event)
         mock_save.assert_not_called()
 
@@ -153,9 +162,9 @@ def test_update_with_message_moves_status_message(
     task_manager: ClientTaskManager, sample_task: Task, sample_message: Message
 ):
     status_message = Message(
-        message_id="status_msg",
+        message_id='status_msg',
         role=Role.agent,
-        parts=[Part(root=TextPart(text="Status"))],
+        parts=[Part(root=TextPart(text='Status'))],
     )
     sample_task.status.message = status_message
     updated_task = task_manager.update_with_message(sample_message, sample_task)