|
3 | 3 |
|
4 | 4 | from unittest.mock import patch |
5 | 5 |
|
6 | | -from a2a.types import DataPart, Part, TextPart |
| 6 | +from a2a.types import ( |
| 7 | + Artifact, |
| 8 | + DataPart, |
| 9 | + Part, |
| 10 | + TextPart, |
| 11 | +) |
7 | 12 | from a2a.utils.artifact import ( |
| 13 | + get_artifact_text, |
8 | 14 | new_artifact, |
9 | 15 | new_data_artifact, |
10 | 16 | new_text_artifact, |
@@ -83,5 +89,71 @@ def test_new_data_artifact_assigns_name_description(self): |
83 | 89 | self.assertEqual(artifact.description, description) |
84 | 90 |
|
85 | 91 |
|
| 92 | +class TestGetArtifactText(unittest.TestCase): |
| 93 | + def test_get_artifact_text_single_part(self): |
| 94 | + # Setup |
| 95 | + artifact = Artifact( |
| 96 | + name='test-artifact', |
| 97 | + parts=[Part(root=TextPart(text='Hello world'))], |
| 98 | + artifact_id='test-artifact-id', |
| 99 | + ) |
| 100 | + |
| 101 | + # Exercise |
| 102 | + result = get_artifact_text(artifact) |
| 103 | + |
| 104 | + # Verify |
| 105 | + assert result == 'Hello world' |
| 106 | + |
| 107 | + def test_get_artifact_text_multiple_parts(self): |
| 108 | + # Setup |
| 109 | + artifact = Artifact( |
| 110 | + name='test-artifact', |
| 111 | + parts=[ |
| 112 | + Part(root=TextPart(text='First line')), |
| 113 | + Part(root=TextPart(text='Second line')), |
| 114 | + Part(root=TextPart(text='Third line')), |
| 115 | + ], |
| 116 | + artifact_id='test-artifact-id', |
| 117 | + ) |
| 118 | + |
| 119 | + # Exercise |
| 120 | + result = get_artifact_text(artifact) |
| 121 | + |
| 122 | + # Verify - default delimiter is newline |
| 123 | + assert result == 'First line\nSecond line\nThird line' |
| 124 | + |
| 125 | + def test_get_artifact_text_custom_delimiter(self): |
| 126 | + # Setup |
| 127 | + artifact = Artifact( |
| 128 | + name='test-artifact', |
| 129 | + parts=[ |
| 130 | + Part(root=TextPart(text='First part')), |
| 131 | + Part(root=TextPart(text='Second part')), |
| 132 | + Part(root=TextPart(text='Third part')), |
| 133 | + ], |
| 134 | + artifact_id='test-artifact-id', |
| 135 | + ) |
| 136 | + |
| 137 | + # Exercise |
| 138 | + result = get_artifact_text(artifact, delimiter=' | ') |
| 139 | + |
| 140 | + # Verify |
| 141 | + assert result == 'First part | Second part | Third part' |
| 142 | + |
| 143 | + def test_get_artifact_text_empty_parts(self): |
| 144 | + # Setup |
| 145 | + artifact = Artifact( |
| 146 | + name='test-artifact', |
| 147 | + parts=[], |
| 148 | + artifact_id='test-artifact-id', |
| 149 | + ) |
| 150 | + |
| 151 | + # Exercise |
| 152 | + result = get_artifact_text(artifact) |
| 153 | + |
| 154 | + # Verify |
| 155 | + assert result == '' |
| 156 | + |
| 157 | + |
86 | 158 | if __name__ == '__main__': |
87 | 159 | unittest.main() |
0 commit comments