imagepipeline/tests/test_ai_modules.py

from __future__ import annotations

import shutil
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

from imagepipeline.core.context import ModuleContext
from imagepipeline.core.exceptions import DependencyError
from imagepipeline.modules.ai_exposure import AIExposureModule
from imagepipeline.modules.ai_tone_map import AIToneMapModule
from imagepipeline.modules.comfy_flux_edit import ComfyFluxEditModule
from imagepipeline.modules.openrouter_edit import OpenRouterEditModule
from imagepipeline.modules.registry import get_module, list_modules
from tests.conftest import make_png

try:
    import numpy  # noqa: F401

    has_numpy = True
except ImportError:
    has_numpy = False

try:
    import torch  # noqa: F401

    has_torch = True
except ImportError:
    has_torch = False

has_magick = bool(shutil.which("magick") or shutil.which("convert"))


class TestAIModuleRegistration:
    def test_ai_modules_registered(self) -> None:
        names = list_modules()
        for name in (
            "ai_exposure",
            "ai_tone_map",
            "openrouter_edit",
            "comfy_flux_edit",
        ):
            assert name in names

    def test_get_ai_modules(self) -> None:
        assert get_module("ai_exposure") is AIExposureModule
        assert get_module("ai_tone_map") is AIToneMapModule
        assert get_module("openrouter_edit") is OpenRouterEditModule
        assert get_module("comfy_flux_edit") is ComfyFluxEditModule


class TestAIParameters:
    def test_ai_exposure_defaults(self) -> None:
        params = AIExposureModule.validate_module_params({})
        assert params["skip_existing"] is True
        assert params["max_edge"] == 2048
        assert params["device"] == "cpu"
        assert params["strength"] == 1.0

    def test_ai_tone_map_defaults(self) -> None:
        params = AIToneMapModule.validate_module_params({})
        assert params["checkpoint"] == ""
        assert params["strength"] == 1.0
        assert params["net_input_size"] == 256

    def test_openrouter_requires_prompt(self) -> None:
        with pytest.raises(ValueError, match="required"):
            OpenRouterEditModule.validate_module_params({})

    def test_openrouter_defaults(self) -> None:
        params = OpenRouterEditModule.validate_module_params({"prompt": "brighten shadows"})
        assert params["model"] == "black-forest-labs/flux.2-klein-4b"
        assert params["strength"] == 0.3
        assert params["api_key_env"] == "OPENROUTER_API_KEY"

    def test_comfy_requires_prompt(self) -> None:
        with pytest.raises(ValueError, match="required"):
            ComfyFluxEditModule.validate_module_params({})


class TestAIToneMapFallback:
    @pytest.mark.skipif(not has_numpy, reason="numpy not installed")
    def test_clahe_fallback_writes_output(self, tmp_path: Path) -> None:
        src = tmp_path / "photo.png"
        make_png(src, width=8, height=8, rgb=(40, 80, 120))
        output_dir = tmp_path / "out"
        output_dir.mkdir()
        ctx = ModuleContext(
            input_paths=[src],
            matched_groups=[],
            output_dir=output_dir,
            params=AIToneMapModule.validate_module_params({"strength": 1.0}),
            pipeline_output_root=tmp_path,
            step_id="ai_tone_map_01",
            logger=None,
        )
        AIToneMapModule().run(ctx)
        dst = output_dir / "photo.png"
        assert dst.is_file()
        assert dst.stat().st_size > 0


class TestSkipExisting:
    @pytest.mark.skipif(not has_numpy, reason="numpy not installed")
    def test_second_run_skips_existing_outputs(self, tmp_path: Path, capsys) -> None:
        src = tmp_path / "photo.png"
        make_png(src, width=8, height=8)
        output_dir = tmp_path / "out"
        output_dir.mkdir()
        params = AIToneMapModule.validate_module_params({"checkpoint": ""})
        from imagepipeline.core.log import PipelineLogger

        logger = PipelineLogger(verbose=True)
        ctx = ModuleContext(
            input_paths=[src],
            matched_groups=[],
            output_dir=output_dir,
            params=params,
            pipeline_output_root=tmp_path,
            step_id="ai_tone_map_01",
            logger=logger,
        )
        module = AIToneMapModule()
        module.run(ctx)
        assert (output_dir / "photo.png").is_file()

        module.run(ctx)
        output = capsys.readouterr().out
        assert "Skipped module ai_tone_map" in output


@pytest.mark.skipif(not has_torch, reason="torch not installed")
class TestAIExposure:
    def test_ai_exposure_processes_image(self, tmp_path: Path) -> None:
        src = tmp_path / "photo.png"
        make_png(src, width=4, height=4, rgb=(30, 60, 90))
        output_dir = tmp_path / "out"
        output_dir.mkdir()

        mock_model = MagicMock()
        mock_device = MagicMock()

        def fake_enhance(_model, image, *, device, strength):
            return image

        with (
            patch.object(AIExposureModule, "_get_model", return_value=(mock_model, mock_device)),
            patch(
                "imagepipeline.ai.zero_dce.enhance_image",
                side_effect=fake_enhance,
            ),
        ):
            ctx = ModuleContext(
                input_paths=[src],
                matched_groups=[],
                output_dir=output_dir,
                params=AIExposureModule.validate_module_params({"max_edge": 0}),
                pipeline_output_root=tmp_path,
                step_id="ai_exposure_01",
                logger=None,
            )
            AIExposureModule().run(ctx)

        assert (output_dir / "photo.png").is_file()


class TestComfyFluxEdit:
    def test_server_unreachable_raises(self, tmp_path: Path) -> None:
        src = tmp_path / "photo.png"
        make_png(src)
        output_dir = tmp_path / "out"
        output_dir.mkdir()
        workflow = tmp_path / "workflow.json"
        workflow.write_text('{"1": {"class_type": "LoadImage", "inputs": {"image": "x"}}}')

        ctx = ModuleContext(
            input_paths=[src],
            matched_groups=[],
            output_dir=output_dir,
            params=ComfyFluxEditModule.validate_module_params(
                {
                    "prompt": "test",
                    "server_url": "http://127.0.0.1:1",
                    "workflow_path": workflow,
                }
            ),
            pipeline_output_root=tmp_path,
            step_id="comfy_flux_edit_01",
            logger=None,
        )
        with pytest.raises(DependencyError, match="not reachable"):
            ComfyFluxEditModule().run(ctx)

    def test_missing_workflow_raises(self, tmp_path: Path) -> None:
        src = tmp_path / "photo.png"
        make_png(src)
        output_dir = tmp_path / "out"
        output_dir.mkdir()
        ctx = ModuleContext(
            input_paths=[src],
            matched_groups=[],
            output_dir=output_dir,
            params=ComfyFluxEditModule.validate_module_params(
                {
                    "prompt": "test",
                    "workflow_path": tmp_path / "missing.json",
                }
            ),
            pipeline_output_root=tmp_path,
            step_id="comfy_flux_edit_01",
            logger=None,
        )
        with pytest.raises(DependencyError, match="workflow not found"):
            ComfyFluxEditModule().run(ctx)


class TestOpenRouterDependencies:
    def test_missing_api_key_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
        monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)
        with pytest.raises(DependencyError, match="OPENROUTER_API_KEY"):
            OpenRouterEditModule.check_dependencies()