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"