From 523117467eec3dd3f27109f0ccc3fe90acd47e5b Mon Sep 17 00:00:00 2001 From: Didier Durand Date: Thu, 15 May 2025 07:26:15 +0200 Subject: [PATCH] adding 13 tests for server/tasks/task_updater.py --- tests/server/tasks/test_task_updater.py | 219 ++++++++++++++++++++++++ 1 file changed, 219 insertions(+) create mode 100644 tests/server/tasks/test_task_updater.py diff --git a/tests/server/tasks/test_task_updater.py b/tests/server/tasks/test_task_updater.py new file mode 100644 index 0000000..9ac632c --- /dev/null +++ b/tests/server/tasks/test_task_updater.py @@ -0,0 +1,219 @@ +import pytest +from unittest.mock import Mock, patch +import uuid + +from a2a.server.events import EventQueue +from a2a.server.tasks import TaskUpdater +from a2a.types import ( + + Message, + Part, + Role, + TaskArtifactUpdateEvent, + TaskState, + TaskStatusUpdateEvent, + TextPart, +) + + +class TestTaskUpdater: + @pytest.fixture + def event_queue(self): + """Create a mock event queue for testing.""" + return Mock(spec=EventQueue) + + @pytest.fixture + def task_updater(self, event_queue): + """Create a TaskUpdater instance for testing.""" + return TaskUpdater( + event_queue=event_queue, + task_id="test-task-id", + context_id="test-context-id" + ) + + @pytest.fixture + def sample_message(self): + """Create a sample message for testing.""" + return Message( + role=Role.agent, + taskId="test-task-id", + contextId="test-context-id", + messageId="test-message-id", + parts=[Part(root=TextPart(text="Test message"))] + ) + + @pytest.fixture + def sample_parts(self): + """Create sample parts for testing.""" + return [Part(root=TextPart(text="Test part"))] + + def test_init(self, event_queue): + """Test that TaskUpdater initializes correctly.""" + task_updater = TaskUpdater( + event_queue=event_queue, + task_id="test-task-id", + context_id="test-context-id" + ) + + assert task_updater.event_queue == event_queue + assert task_updater.task_id == "test-task-id" + assert task_updater.context_id == "test-context-id" + + def test_update_status_without_message(self, task_updater, event_queue): + """Test updating status without a message.""" + task_updater.update_status(TaskState.working) + + event_queue.enqueue_event.assert_called_once() + event = event_queue.enqueue_event.call_args[0][0] + + assert isinstance(event, TaskStatusUpdateEvent) + assert event.taskId == "test-task-id" + assert event.contextId == "test-context-id" + assert event.final is False + assert event.status.state == TaskState.working + assert event.status.message is None + + def test_update_status_with_message(self, task_updater, event_queue, sample_message): + """Test updating status with a message.""" + task_updater.update_status(TaskState.working, message=sample_message) + + event_queue.enqueue_event.assert_called_once() + event = event_queue.enqueue_event.call_args[0][0] + + assert isinstance(event, TaskStatusUpdateEvent) + assert event.taskId == "test-task-id" + assert event.contextId == "test-context-id" + assert event.final is False + assert event.status.state == TaskState.working + assert event.status.message == sample_message + + def test_update_status_final(self, task_updater, event_queue): + """Test updating status with final=True.""" + task_updater.update_status(TaskState.completed, final=True) + + event_queue.enqueue_event.assert_called_once() + event = event_queue.enqueue_event.call_args[0][0] + + assert isinstance(event, TaskStatusUpdateEvent) + assert event.final is True + assert event.status.state == TaskState.completed + + def test_add_artifact_with_custom_id_and_name(self, task_updater, event_queue, sample_parts): + """Test adding an artifact with a custom ID and name.""" + task_updater.add_artifact( + parts=sample_parts, + artifact_id="custom-artifact-id", + name="Custom Artifact" + ) + + event_queue.enqueue_event.assert_called_once() + event = event_queue.enqueue_event.call_args[0][0] + + assert isinstance(event, TaskArtifactUpdateEvent) + assert event.artifact.artifactId == "custom-artifact-id" + assert event.artifact.name == "Custom Artifact" + assert event.artifact.parts == sample_parts + + def test_complete_without_message(self, task_updater, event_queue): + """Test marking a task as completed without a message.""" + task_updater.complete() + + event_queue.enqueue_event.assert_called_once() + event = event_queue.enqueue_event.call_args[0][0] + + assert isinstance(event, TaskStatusUpdateEvent) + assert event.status.state == TaskState.completed + assert event.final is True + assert event.status.message is None + + def test_complete_with_message(self, task_updater, event_queue, sample_message): + """Test marking a task as completed with a message.""" + task_updater.complete(message=sample_message) + + event_queue.enqueue_event.assert_called_once() + event = event_queue.enqueue_event.call_args[0][0] + + assert isinstance(event, TaskStatusUpdateEvent) + assert event.status.state == TaskState.completed + assert event.final is True + assert event.status.message == sample_message + + def test_submit_without_message(self, task_updater, event_queue): + """Test marking a task as submitted without a message.""" + task_updater.submit() + + event_queue.enqueue_event.assert_called_once() + event = event_queue.enqueue_event.call_args[0][0] + + assert isinstance(event, TaskStatusUpdateEvent) + assert event.status.state == TaskState.submitted + assert event.final is False + assert event.status.message is None + + def test_submit_with_message(self, task_updater, event_queue, sample_message): + """Test marking a task as submitted with a message.""" + task_updater.submit(message=sample_message) + + event_queue.enqueue_event.assert_called_once() + event = event_queue.enqueue_event.call_args[0][0] + + assert isinstance(event, TaskStatusUpdateEvent) + assert event.status.state == TaskState.submitted + assert event.final is False + assert event.status.message == sample_message + + def test_start_work_without_message(self, task_updater, event_queue): + """Test marking a task as working without a message.""" + task_updater.start_work() + + event_queue.enqueue_event.assert_called_once() + event = event_queue.enqueue_event.call_args[0][0] + + assert isinstance(event, TaskStatusUpdateEvent) + assert event.status.state == TaskState.working + assert event.final is False + assert event.status.message is None + + def test_start_work_with_message(self, task_updater, event_queue, sample_message): + """Test marking a task as working with a message.""" + task_updater.start_work(message=sample_message) + + event_queue.enqueue_event.assert_called_once() + event = event_queue.enqueue_event.call_args[0][0] + + assert isinstance(event, TaskStatusUpdateEvent) + assert event.status.state == TaskState.working + assert event.final is False + assert event.status.message == sample_message + + def test_new_agent_message(self, task_updater, sample_parts): + """Test creating a new agent message.""" + with patch('uuid.uuid4', return_value=uuid.UUID('12345678-1234-5678-1234-567812345678')): + message = task_updater.new_agent_message(parts=sample_parts) + + assert message.role == Role.agent + assert message.taskId == "test-task-id" + assert message.contextId == "test-context-id" + assert message.messageId == "12345678-1234-5678-1234-567812345678" + assert message.parts == sample_parts + assert message.final is False + assert message.metadata is None + + def test_new_agent_message_with_metadata_and_final(self, task_updater, sample_parts): + """Test creating a new agent message with metadata and final=True.""" + metadata = {"key": "value"} + + with patch('uuid.uuid4', return_value=uuid.UUID('12345678-1234-5678-1234-567812345678')): + message = task_updater.new_agent_message( + parts=sample_parts, + final=True, + metadata=metadata + ) + + assert message.role == Role.agent + assert message.taskId == "test-task-id" + assert message.contextId == "test-context-id" + assert message.messageId == "12345678-1234-5678-1234-567812345678" + assert message.parts == sample_parts + assert message.final is True + assert message.metadata == metadata