Add unit tests for ArtifactStreamer helper

Author: kevinlu310

tests/utils/test_artifact.py

@@ -7,9 +7,11 @@
     Artifact,
     DataPart,
     Part,
+    TaskArtifactUpdateEvent,
     TextPart,
 )
 from a2a.utils.artifact import (
+    ArtifactStreamer,
     get_artifact_text,
     new_artifact,
     new_data_artifact,
@@ -155,5 +157,105 @@ def test_get_artifact_text_empty_parts(self):
         assert result == ''
 
 
+class TestArtifactStreamer(unittest.TestCase):
+    def setUp(self):
+        self.context_id = 'ctx-123'
+        self.task_id = 'task-456'
+
+    def test_generates_stable_artifact_id(self):
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        e1 = streamer.append('hello ')
+        e2 = streamer.append('world')
+        self.assertEqual(e1.artifact.artifact_id, e2.artifact.artifact_id)
+
+    def test_uses_explicit_artifact_id(self):
+        streamer = ArtifactStreamer(
+            self.context_id, self.task_id, artifact_id='my-fixed-id'
+        )
+        event = streamer.append('chunk')
+        self.assertEqual(event.artifact.artifact_id, 'my-fixed-id')
+
+    @patch('a2a.utils.artifact.uuid.uuid4')
+    def test_generated_id_comes_from_uuid4(self, mock_uuid4):
+        mock_uuid = uuid.UUID('abcdef12-1234-5678-1234-567812345678')
+        mock_uuid4.return_value = mock_uuid
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        self.assertEqual(streamer.artifact_id, str(mock_uuid))
+
+    def test_default_name_is_response(self):
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        event = streamer.append('text')
+        self.assertEqual(event.artifact.name, 'response')
+
+    def test_custom_name(self):
+        streamer = ArtifactStreamer(
+            self.context_id, self.task_id, name='summary'
+        )
+        event = streamer.append('text')
+        self.assertEqual(event.artifact.name, 'summary')
+
+    def test_append_returns_task_artifact_update_event(self):
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        event = streamer.append('chunk')
+        self.assertIsInstance(event, TaskArtifactUpdateEvent)
+
+    def test_append_sets_correct_context_and_task(self):
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        event = streamer.append('chunk')
+        self.assertEqual(event.context_id, self.context_id)
+        self.assertEqual(event.task_id, self.task_id)
+
+    def test_append_sets_append_true_last_chunk_false(self):
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        event = streamer.append('chunk')
+        self.assertTrue(event.append)
+        self.assertFalse(event.last_chunk)
+
+    def test_append_creates_single_text_part(self):
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        event = streamer.append('hello')
+        self.assertEqual(len(event.artifact.parts), 1)
+        self.assertIsInstance(event.artifact.parts[0].root, TextPart)
+        self.assertEqual(event.artifact.parts[0].root.text, 'hello')
+
+    def test_finalize_returns_task_artifact_update_event(self):
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        event = streamer.finalize()
+        self.assertIsInstance(event, TaskArtifactUpdateEvent)
+
+    def test_finalize_sets_append_true_last_chunk_true(self):
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        event = streamer.finalize()
+        self.assertTrue(event.append)
+        self.assertTrue(event.last_chunk)
+
+    def test_finalize_has_empty_parts(self):
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        event = streamer.finalize()
+        self.assertEqual(event.artifact.parts, [])
+
+    def test_finalize_uses_same_artifact_id_as_append(self):
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        append_event = streamer.append('text')
+        finalize_event = streamer.finalize()
+        self.assertEqual(
+            append_event.artifact.artifact_id,
+            finalize_event.artifact.artifact_id,
+        )
+
+    def test_multiple_appends_all_share_artifact_id(self):
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        events = [streamer.append(f'chunk-{i}') for i in range(5)]
+        ids = {e.artifact.artifact_id for e in events}
+        self.assertEqual(len(ids), 1)
+
+    def test_multiple_appends_carry_distinct_text(self):
+        streamer = ArtifactStreamer(self.context_id, self.task_id)
+        texts = ['Hello, ', 'world', '!']
+        events = [streamer.append(t) for t in texts]
+        result_texts = [e.artifact.parts[0].root.text for e in events]
+        self.assertEqual(result_texts, texts)
+
+
 if __name__ == '__main__':
     unittest.main()