110 lines
3.9 KiB
Python
110 lines
3.9 KiB
Python
from pathlib import Path
|
|
from subprocess import CompletedProcess
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from app.config import Settings
|
|
from app.engines.base import UpscaleEngineUnavailable
|
|
from app.image_io import DownloadedImage
|
|
from app.upscaler import build_upscaler
|
|
|
|
|
|
def make_downloaded_image(tmp_path: Path) -> DownloadedImage:
|
|
source = tmp_path / "source.png"
|
|
Image.new("RGBA", (12, 8), (25, 50, 75, 255)).save(source, "PNG")
|
|
|
|
return DownloadedImage(
|
|
path=source,
|
|
width=12,
|
|
height=8,
|
|
mime="image/png",
|
|
filesize=source.stat().st_size,
|
|
)
|
|
|
|
|
|
def make_runtime_settings(tmp_path: Path) -> Settings:
|
|
binary = tmp_path / "bin" / "realesrgan-ncnn-vulkan"
|
|
binary.parent.mkdir(parents=True)
|
|
binary.write_text("#!/bin/sh\nexit 0\n", encoding="utf-8")
|
|
binary.chmod(0o755)
|
|
|
|
model_dir = tmp_path / "models"
|
|
model_dir.mkdir(parents=True)
|
|
(model_dir / "realesrgan-x4plus.param").write_text("param", encoding="utf-8")
|
|
(model_dir / "realesrgan-x4plus.bin").write_text("bin", encoding="utf-8")
|
|
|
|
return Settings(
|
|
engine="realesrgan-ncnn",
|
|
device="vulkan",
|
|
token="secret-token",
|
|
tmp_dir=str(tmp_path / "tmp"),
|
|
output_dir=str(tmp_path / "output"),
|
|
realesrgan_bin=str(binary),
|
|
realesrgan_model_dir=str(model_dir),
|
|
realesrgan_anime_model="realesrgan-x4plus-anime",
|
|
realesrgan_allow_model_fallback=True,
|
|
)
|
|
|
|
|
|
def test_realesrgan_command_is_built_without_shell_and_2x_records_downsample(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
settings = make_runtime_settings(tmp_path)
|
|
engine = build_upscaler(settings)
|
|
downloaded = make_downloaded_image(tmp_path)
|
|
captured: dict[str, object] = {}
|
|
|
|
def fake_run(command, **kwargs):
|
|
captured["command"] = command
|
|
captured["kwargs"] = kwargs
|
|
output_index = command.index("-o") + 1
|
|
output_path = Path(command[output_index])
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
Image.new("RGBA", (downloaded.width * 4, downloaded.height * 4), (120, 80, 20, 255)).save(output_path, "PNG")
|
|
|
|
return CompletedProcess(command, 0, stdout="ok", stderr="")
|
|
|
|
monkeypatch.setattr("app.engines.realesrgan_ncnn_engine.subprocess.run", fake_run)
|
|
|
|
result = engine.upscale(downloaded, 2, "illustration", "webp")
|
|
|
|
assert result.image.size == (downloaded.width * 2, downloaded.height * 2)
|
|
assert result.metadata["requested_scale"] == 2
|
|
assert result.metadata["native_model_scale"] == 4
|
|
assert result.metadata["post_downsampled"] is True
|
|
assert result.metadata["model_fallback"] is True
|
|
assert result.metadata["used_model"] == "realesrgan-x4plus"
|
|
assert captured["command"][0] == settings.realesrgan_bin
|
|
assert captured["kwargs"].get("check") is False
|
|
assert captured["kwargs"].get("shell", False) is False
|
|
|
|
|
|
def test_missing_binary_raises_safe_engine_unavailable(tmp_path: Path) -> None:
|
|
engine = build_upscaler(
|
|
Settings(
|
|
engine="realesrgan-ncnn",
|
|
token="secret-token",
|
|
tmp_dir=str(tmp_path / "tmp"),
|
|
output_dir=str(tmp_path / "output"),
|
|
realesrgan_bin=str(tmp_path / "bin" / "missing-binary"),
|
|
realesrgan_model_dir=str(tmp_path / "models"),
|
|
)
|
|
)
|
|
|
|
with pytest.raises(UpscaleEngineUnavailable, match="Upscale engine is not available"):
|
|
engine.upscale(make_downloaded_image(tmp_path), 4, "artwork", "webp")
|
|
|
|
|
|
def test_pillow_engine_still_works(tmp_path: Path) -> None:
|
|
engine = build_upscaler(
|
|
Settings(
|
|
engine="pillow",
|
|
token="secret-token",
|
|
tmp_dir=str(tmp_path / "tmp"),
|
|
output_dir=str(tmp_path / "output"),
|
|
)
|
|
)
|
|
|
|
result = engine.upscale(make_downloaded_image(tmp_path), 2, "standard", "webp")
|
|
|
|
assert result.metadata["engine"] == "pillow"
|
|
assert result.metadata["real_ai_upscale"] is False |