|
7 | 7 | Artifact, |
8 | 8 | DataPart, |
9 | 9 | Part, |
| 10 | + TaskArtifactUpdateEvent, |
10 | 11 | TextPart, |
11 | 12 | ) |
12 | 13 | from a2a.utils.artifact import ( |
| 14 | + ArtifactStreamer, |
13 | 15 | get_artifact_text, |
14 | 16 | new_artifact, |
15 | 17 | new_data_artifact, |
@@ -155,5 +157,131 @@ def test_get_artifact_text_empty_parts(self): |
155 | 157 | assert result == '' |
156 | 158 |
|
157 | 159 |
|
| 160 | +class TestArtifactStreamer(unittest.TestCase): |
| 161 | + def setUp(self): |
| 162 | + self.context_id = 'ctx-123' |
| 163 | + self.task_id = 'task-456' |
| 164 | + self.name = 'response' |
| 165 | + |
| 166 | + def test_stable_artifact_id_across_appends(self): |
| 167 | + streamer = ArtifactStreamer( |
| 168 | + self.context_id, self.task_id, name=self.name |
| 169 | + ) |
| 170 | + event1 = streamer.append('Hello ') |
| 171 | + event2 = streamer.append('world') |
| 172 | + self.assertEqual( |
| 173 | + event1.artifact.artifact_id, event2.artifact.artifact_id |
| 174 | + ) |
| 175 | + |
| 176 | + def test_append_returns_correct_event_type(self): |
| 177 | + streamer = ArtifactStreamer( |
| 178 | + self.context_id, self.task_id, name=self.name |
| 179 | + ) |
| 180 | + event = streamer.append('chunk') |
| 181 | + self.assertIsInstance(event, TaskArtifactUpdateEvent) |
| 182 | + |
| 183 | + def test_append_sets_append_true_last_chunk_false(self): |
| 184 | + streamer = ArtifactStreamer( |
| 185 | + self.context_id, self.task_id, name=self.name |
| 186 | + ) |
| 187 | + event = streamer.append('chunk') |
| 188 | + self.assertTrue(event.append) |
| 189 | + self.assertFalse(event.last_chunk) |
| 190 | + |
| 191 | + def test_append_sets_context_and_task_ids(self): |
| 192 | + streamer = ArtifactStreamer( |
| 193 | + self.context_id, self.task_id, name=self.name |
| 194 | + ) |
| 195 | + event = streamer.append('chunk') |
| 196 | + self.assertEqual(event.context_id, self.context_id) |
| 197 | + self.assertEqual(event.task_id, self.task_id) |
| 198 | + |
| 199 | + def test_append_sets_text_content(self): |
| 200 | + streamer = ArtifactStreamer( |
| 201 | + self.context_id, self.task_id, name=self.name |
| 202 | + ) |
| 203 | + event = streamer.append('Hello world') |
| 204 | + self.assertEqual(len(event.artifact.parts), 1) |
| 205 | + self.assertEqual(event.artifact.parts[0].root.text, 'Hello world') |
| 206 | + |
| 207 | + def test_append_sets_artifact_name_and_description(self): |
| 208 | + streamer = ArtifactStreamer( |
| 209 | + self.context_id, |
| 210 | + self.task_id, |
| 211 | + name='my-artifact', |
| 212 | + description='A streamed response', |
| 213 | + ) |
| 214 | + event = streamer.append('chunk') |
| 215 | + self.assertEqual(event.artifact.name, 'my-artifact') |
| 216 | + self.assertEqual(event.artifact.description, 'A streamed response') |
| 217 | + |
| 218 | + def test_finalize_sets_last_chunk_true(self): |
| 219 | + streamer = ArtifactStreamer( |
| 220 | + self.context_id, self.task_id, name=self.name |
| 221 | + ) |
| 222 | + event = streamer.finalize('done') |
| 223 | + self.assertTrue(event.append) |
| 224 | + self.assertTrue(event.last_chunk) |
| 225 | + |
| 226 | + def test_finalize_with_empty_text(self): |
| 227 | + streamer = ArtifactStreamer( |
| 228 | + self.context_id, self.task_id, name=self.name |
| 229 | + ) |
| 230 | + event = streamer.finalize() |
| 231 | + self.assertEqual(event.artifact.parts[0].root.text, '') |
| 232 | + self.assertTrue(event.last_chunk) |
| 233 | + |
| 234 | + def test_finalize_uses_same_artifact_id(self): |
| 235 | + streamer = ArtifactStreamer( |
| 236 | + self.context_id, self.task_id, name=self.name |
| 237 | + ) |
| 238 | + append_event = streamer.append('chunk') |
| 239 | + finalize_event = streamer.finalize() |
| 240 | + self.assertEqual( |
| 241 | + append_event.artifact.artifact_id, |
| 242 | + finalize_event.artifact.artifact_id, |
| 243 | + ) |
| 244 | + |
| 245 | + def test_append_after_finalize_raises(self): |
| 246 | + streamer = ArtifactStreamer( |
| 247 | + self.context_id, self.task_id, name=self.name |
| 248 | + ) |
| 249 | + streamer.finalize() |
| 250 | + with self.assertRaises(RuntimeError): |
| 251 | + streamer.append('too late') |
| 252 | + |
| 253 | + def test_double_finalize_raises(self): |
| 254 | + streamer = ArtifactStreamer( |
| 255 | + self.context_id, self.task_id, name=self.name |
| 256 | + ) |
| 257 | + streamer.finalize() |
| 258 | + with self.assertRaises(RuntimeError): |
| 259 | + streamer.finalize() |
| 260 | + |
| 261 | + def test_artifact_id_property(self): |
| 262 | + streamer = ArtifactStreamer( |
| 263 | + self.context_id, self.task_id, name=self.name |
| 264 | + ) |
| 265 | + artifact_id = streamer.artifact_id |
| 266 | + self.assertIsInstance(artifact_id, str) |
| 267 | + self.assertTrue(len(artifact_id) > 0) |
| 268 | + |
| 269 | + @patch('uuid.uuid4') |
| 270 | + def test_artifact_id_from_uuid(self, mock_uuid4): |
| 271 | + mock_uuid = uuid.UUID('12345678-1234-5678-1234-567812345678') |
| 272 | + mock_uuid4.return_value = mock_uuid |
| 273 | + streamer = ArtifactStreamer( |
| 274 | + self.context_id, self.task_id, name=self.name |
| 275 | + ) |
| 276 | + self.assertEqual(streamer.artifact_id, str(mock_uuid)) |
| 277 | + |
| 278 | + def test_description_defaults_to_none(self): |
| 279 | + streamer = ArtifactStreamer( |
| 280 | + self.context_id, self.task_id, name=self.name |
| 281 | + ) |
| 282 | + event = streamer.append('chunk') |
| 283 | + self.assertIsNone(event.artifact.description) |
| 284 | + |
| 285 | + |
158 | 286 | if __name__ == '__main__': |
159 | 287 | unittest.main() |
0 commit comments