1 files changed, 34 insertions, 0 deletions
diff --git a/tests/test_images.py b/tests/test_images.py
new file mode 100644
index 0000000..f2ac44e
--- /dev/null
+++ b/tests/test_images.py
@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import base64
+import io
+
+import pytest
+from PIL import Image
+
+from mcp_sd15.errors import InvalidImageError
+from mcp_sd15.images import decode_image, encode_png
+
+
+def test_decodes_raw_base64_and_data_uri(image_base64: str) -> None:
+    assert decode_image(image_base64).mode == "RGB"
+    assert decode_image(f"data:image/png;base64,{image_base64}", mask=True).mode == "L"
+
+
+@pytest.mark.parametrize("value", ["not base64!", "data:text/plain;base64,YQ==", ""])
+def test_rejects_invalid_image_input(value: str) -> None:
+    with pytest.raises(InvalidImageError):
+        decode_image(value)
+
+
+def test_rejects_non_image_base64() -> None:
+    value = base64.b64encode(b"plain text").decode()
+    with pytest.raises(InvalidImageError):
+        decode_image(value)
+
+
+def test_png_output_contains_metadata() -> None:
+    data = encode_png(Image.new("RGB", (2, 2)), {"seed": 42, "mode": "test"})
+    with Image.open(io.BytesIO(data)) as result:
+        assert result.format == "PNG"
+        assert result.info["seed"] == "42"