1 files changed, 118 insertions, 0 deletions
diff --git a/tests/test_service.py b/tests/test_service.py
new file mode 100644
index 0000000..da82204
--- /dev/null
+++ b/tests/test_service.py
@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import asyncio
+import io
+import threading
+import time
+from dataclasses import replace
+
+import pytest
+from PIL import Image
+
+from mcp_sd15.config import ControlModelConfig, ControlNetConfig
+from mcp_sd15.errors import GenerationError, InvalidRequestError
+from mcp_sd15.schema import GuidedImageRequest, InpaintImageRequest, TextToImageRequest
+from mcp_sd15.service import ImageGenerationService
+
+
+class FakeBackend:
+    def __init__(self, delay: float = 0) -> None:
+        self.calls: list[str] = []
+        self.delay = delay
+        self.active = 0
+        self.max_active = 0
+        self.guard = threading.Lock()
+
+    def _result(self, name: str) -> Image.Image:
+        with self.guard:
+            self.active += 1
+            self.max_active = max(self.max_active, self.active)
+        time.sleep(self.delay)
+        with self.guard:
+            self.active -= 1
+        self.calls.append(name)
+        return Image.new("RGB", (4, 4), "red")
+
+    def text_to_image(self, request: object, seed: int) -> Image.Image:
+        return self._result("text")
+
+    def guided_image(self, request: object, image: Image.Image, seed: int) -> Image.Image:
+        assert image.mode == "RGB"
+        return self._result("guided")
+
+    def inpaint_image(
+        self, request: object, image: Image.Image, mask: Image.Image, seed: int
+    ) -> Image.Image:
+        assert mask.mode == "L"
+        return self._result("inpaint")
+
+
+async def test_routes_all_generation_types(app_config, image_base64: str) -> None:
+    backend = FakeBackend()
+    service = ImageGenerationService(app_config, backend)
+
+    text = await service.generate_image(TextToImageRequest.create("cat"))
+    guided = await service.generate_guided_image(
+        GuidedImageRequest.create("cat", "depth", image_base64)
+    )
+    inpaint = await service.inpaint_image(
+        InpaintImageRequest.create("cat", image_base64, image_base64)
+    )
+
+    assert backend.calls == ["text", "guided", "inpaint"]
+    for result in (text, guided, inpaint):
+        assert result.png_data.startswith(b"\x89PNG")
+        assert isinstance(result.metadata["seed"], int)
+
+
+async def test_generation_is_strictly_serialized(app_config) -> None:
+    backend = FakeBackend(delay=0.03)
+    service = ImageGenerationService(app_config, backend)
+    await asyncio.gather(
+        *(service.generate_image(TextToImageRequest.create(f"image {i}")) for i in range(4))
+    )
+    assert backend.max_active == 1
+
+
+async def test_disabled_controlnet_fails_before_backend(app_config, image_base64: str) -> None:
+    disabled = replace(
+        app_config,
+        controlnet=ControlNetConfig(
+            depth=ControlModelConfig(False, None, "midas"),
+            openpose=app_config.controlnet.openpose,
+            inpaint=app_config.controlnet.inpaint,
+        ),
+    )
+    service = ImageGenerationService(disabled, FakeBackend())
+    with pytest.raises(InvalidRequestError, match="disabled"):
+        await service.generate_guided_image(GuidedImageRequest.create("cat", "depth", image_base64))
+
+
+async def test_backend_error_is_sanitized(app_config) -> None:
+    class BrokenBackend(FakeBackend):
+        def text_to_image(self, request: object, seed: int) -> Image.Image:
+            raise RuntimeError("boom")
+
+    service = ImageGenerationService(app_config, BrokenBackend())
+    with pytest.raises(GenerationError, match="boom"):
+        await service.generate_image(TextToImageRequest.create("cat"))
+
+
+async def test_cuda_oom_is_classified(app_config) -> None:
+    class OomBackend(FakeBackend):
+        def text_to_image(self, request: object, seed: int) -> Image.Image:
+            raise RuntimeError("CUDA out of memory while allocating")
+
+    service = ImageGenerationService(app_config, OomBackend())
+    with pytest.raises(GenerationError, match="CUDA out of memory"):
+        await service.generate_image(TextToImageRequest.create("cat"))
+
+
+async def test_optional_output_is_saved(app_config) -> None:
+    enabled = replace(app_config, output=replace(app_config.output, enabled=True))
+    service = ImageGenerationService(enabled, FakeBackend())
+    await service.generate_image(TextToImageRequest.create("cat"))
+    files = list(enabled.output.directory.glob("*.png"))
+    assert len(files) == 1
+    with Image.open(io.BytesIO(files[0].read_bytes())) as image:
+        assert image.format == "PNG"