EvoScientist/tests/test_image_gen_service.py

"""Tests for image_gen.service."""

from pathlib import Path

import pytest

from EvoScientist.image_gen import service
from EvoScientist.image_gen.adapters.base import ImageGenError

PNG_BYTES = b"\x89PNG\r\n\x1a\n" + b"svc"


class _FakeAdapter:
    """Records calls; returns fixed bytes."""

    instances: list["_FakeAdapter"] = []

    def __init__(self, entry, *, client=None, timeout=120.0):
        self.entry = entry
        self.calls: list[dict] = []
        _FakeAdapter.instances.append(self)

    async def generate(self, **kwargs):
        self.calls.append({"kind": "generate", **kwargs})
        return [PNG_BYTES]

    async def edit(self, **kwargs):
        self.calls.append({"kind": "edit", **kwargs})
        return [PNG_BYTES]


@pytest.fixture(autouse=True)
def fake_adapters(monkeypatch, tmp_path):
    _FakeAdapter.instances = []
    monkeypatch.setattr(
        service, "ADAPTER_CLASSES", {"openai": _FakeAdapter, "gemini": _FakeAdapter}
    )
    config = tmp_path / "config.yaml"
    config.write_text(
        """
image_generation:
  default_model: gpt-image-2
  models:
    - id: gpt-image-2
      name: GPT Image 2
      provider: openai
      api_key: sk-x
    - id: imagen-4.0-generate-001
      provider: gemini
      api_key: sk-y
      supports_edit: false
""",
        encoding="utf-8",
    )
    monkeypatch.setattr(service, "_config_path", lambda: config)
    return config


async def test_generate_saves_to_artifacts(tmp_path):
    result = await service.generate_for_workspace(tmp_path, prompt="a cat")
    assert result["ok"] is True
    assert result["model"] == "gpt-image-2"
    assert result["paths"][0].startswith("artifacts/generated_")
    saved = tmp_path / result["paths"][0]
    assert saved.read_bytes() == PNG_BYTES
    adapter = _FakeAdapter.instances[0]
    assert adapter.entry.id == "gpt-image-2"
    assert adapter.calls[0]["prompt"] == "a cat"


async def test_generate_explicit_output_path(tmp_path):
    result = await service.generate_for_workspace(
        tmp_path, prompt="x", output_path="artifacts/cover.png"
    )
    assert result["paths"] == ["artifacts/cover.png"]
    assert (tmp_path / "artifacts" / "cover.png").is_file()


async def test_output_path_traversal_rejected(tmp_path):
    with pytest.raises(ImageGenError):
        await service.generate_for_workspace(
            tmp_path, prompt="x", output_path="../escape.png"
        )
    with pytest.raises(ImageGenError):
        await service.generate_for_workspace(
            tmp_path, prompt="x", output_path="/abs/path.png"
        )


async def test_output_path_must_be_artifacts_png(tmp_path):
    with pytest.raises(ImageGenError):
        await service.generate_for_workspace(
            tmp_path, prompt="x", output_path="other/cover.png"
        )
    with pytest.raises(ImageGenError):
        await service.generate_for_workspace(
            tmp_path, prompt="x", output_path="artifacts/cover.jpg"
        )


async def test_unknown_model_lists_available(tmp_path):
    with pytest.raises(ImageGenError) as excinfo:
        await service.generate_for_workspace(tmp_path, prompt="x", model="nope")
    message = str(excinfo.value)
    assert "GPT Image 2" in message and "imagen-4.0-generate-001" in message


async def test_no_models_configured(tmp_path, monkeypatch):
    empty = tmp_path / "empty.yaml"
    empty.write_text("webui_port: 4716\n", encoding="utf-8")
    monkeypatch.setattr(service, "_config_path", lambda: empty)
    with pytest.raises(ImageGenError, match="no image models"):
        await service.generate_for_workspace(tmp_path, prompt="x")


async def test_n_greater_than_one_numbered(tmp_path):
    class _TwoAdapter(_FakeAdapter):
        async def generate(self, **kwargs):
            return [PNG_BYTES, PNG_BYTES]

    monkeypatch = pytest.MonkeyPatch()
    monkeypatch.setattr(
        service, "ADAPTER_CLASSES", {"openai": _TwoAdapter, "gemini": _TwoAdapter}
    )
    result = await service.generate_for_workspace(
        tmp_path, prompt="x", output_path="artifacts/v.png", n=2
    )
    assert result["paths"] == ["artifacts/v_1.png", "artifacts/v_2.png"]
    monkeypatch.undo()


async def test_edit_reads_input_and_saves(tmp_path):
    src = tmp_path / "artifacts" / "src.png"
    src.parent.mkdir(parents=True)
    src.write_bytes(PNG_BYTES)
    result = await service.edit_for_workspace(
        tmp_path, image_path="artifacts/src.png", prompt="make blue"
    )
    assert result["ok"] is True
    assert result["paths"][0].startswith("artifacts/edited_")
    adapter = _FakeAdapter.instances[0]
    assert adapter.calls[0]["kind"] == "edit"


async def test_edit_rejects_bad_input(tmp_path):
    with pytest.raises(ImageGenError):
        await service.edit_for_workspace(
            tmp_path, image_path="artifacts/missing.png", prompt="x"
        )
    outside = tmp_path / "outside.txt"
    outside.write_text("hi", encoding="utf-8")
    with pytest.raises(ImageGenError):
        await service.edit_for_workspace(
            tmp_path, image_path="outside.txt", prompt="x"
        )


async def test_edit_blocked_when_model_lacks_support(tmp_path):
    with pytest.raises(ImageGenError, match="edit"):
        await service.edit_for_workspace(
            tmp_path,
            image_path="artifacts/whatever.png",
            prompt="x",
            model="imagen-4.0-generate-001",
        )