| 1 | import base64 |
| 2 | from types import SimpleNamespace |
| 3 | import unittest |
| 4 | from unittest.mock import patch |
| 5 | |
| 6 | from PIL import Image |
| 7 | |
| 8 | from tools.image_response import image_from_response_part |
| 9 | |
| 10 | |
| 11 | class ImageResponseTests(unittest.TestCase): |
| 12 | def test_extracts_image_when_part_has_no_as_image_method(self): |
| 13 | expected = Image.new("RGB", (1, 1), (255, 0, 0)) |
| 14 | with patch("tools.image_response.Image.open", return_value=expected): |
| 15 | part = SimpleNamespace(inline_data=SimpleNamespace(data=b"fake-png")) |
| 16 | image = image_from_response_part(part) |
| 17 | self.assertEqual(image.size, (1, 1)) |
| 18 | |
| 19 | def test_extracts_base64_data_url(self): |
| 20 | expected = Image.new("RGB", (1, 1), (255, 0, 0)) |
| 21 | payload = "data:image/png;base64," + base64.b64encode(b"fake-png").decode("ascii") |
| 22 | with patch("tools.image_response.Image.open", return_value=expected): |
| 23 | part = {"inline_data": {"data": payload}} |
| 24 | image = image_from_response_part(part) |
| 25 | self.assertEqual(image.size, (1, 1)) |
| 26 | |
| 27 | |
| 28 | if __name__ == "__main__": |
| 29 | unittest.main() |
| 30 |