wehub-resource-sync
30 lines
1.1 KiB
Python
30 lines
1.1 KiB
Python
import base64
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from PIL import Image
|
|
|
|
from tools.image_response import image_from_response_part
|
|
|
|
|
|
class ImageResponseTests(unittest.TestCase):
|
|
def test_extracts_image_when_part_has_no_as_image_method(self):
|
|
expected = Image.new("RGB", (1, 1), (255, 0, 0))
|
|
with patch("tools.image_response.Image.open", return_value=expected):
|
|
part = SimpleNamespace(inline_data=SimpleNamespace(data=b"fake-png"))
|
|
image = image_from_response_part(part)
|
|
self.assertEqual(image.size, (1, 1))
|
|
|
|
def test_extracts_base64_data_url(self):
|
|
expected = Image.new("RGB", (1, 1), (255, 0, 0))
|
|
payload = "data:image/png;base64," + base64.b64encode(b"fake-png").decode("ascii")
|
|
with patch("tools.image_response.Image.open", return_value=expected):
|
|
part = {"inline_data": {"data": payload}}
|
|
image = image_from_response_part(part)
|
|
self.assertEqual(image.size, (1, 1))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|