tests/test_image_upscale_progress.py

from __future__ import annotations

from pathlib import Path

from PIL import Image
import pytest
import torch

import kneiff.utils.image.upscale as image_upscale
from kneiff.progress import ProgressUpdate


class _FakeUpscaleModel:
    """Small runtime descriptor used without loading model weights."""

    device = torch.device("cpu")
    dtype = torch.float32
    scale = 4


def _install_fake_upscale_runtime(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(
        image_upscale,
        "pick_device",
        lambda gpu, cpu: torch.device("cpu"),
    )
    monkeypatch.setattr(
        image_upscale,
        "resolve_weights_filename",
        lambda **kwargs: "4x-test.safetensors",
    )
    monkeypatch.setattr(
        image_upscale,
        "download_weights",
        lambda **kwargs: Path("cached-weights.safetensors"),
    )
    monkeypatch.setattr(
        image_upscale,
        "load_sr_model",
        lambda **kwargs: _FakeUpscaleModel(),
    )
    monkeypatch.setattr(
        image_upscale,
        "upscale_pil_image",
        lambda **kwargs: Image.new("RGB", (8, 8), color="green"),
    )


def test_local_upscale_reports_loading_phases_and_late_image_total(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    _install_fake_upscale_runtime(monkeypatch)
    first = tmp_path / "first.png"
    second = tmp_path / "second.jpg"
    Image.new("RGB", (2, 2), color="red").save(first)
    Image.new("RGB", (2, 2), color="blue").save(second)
    updates: list[ProgressUpdate] = []

    result = image_upscale.run_local_upscale(
        image_upscale.LocalUpscaleRequest(
            inputs=[first, second],
            out_dir=tmp_path / "output",
            cpu=True,
        ),
        progress_callback=updates.append,
    )

    assert [update.description for update in updates[:4]] == [
        "Selecting upscale device",
        "Resolving model weights",
        "Downloading model weights",
        "Loading upscale model",
    ]
    assert sum(update.additional_total for update in updates) == 2
    assert sum(update.advance for update in updates) == 2
    assert len(result.results) == 2
    assert all(item.output_path.is_file() for item in result.results)


def test_local_upscale_main_preserves_legacy_redirected_output(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
    capsys: pytest.CaptureFixture[str],
) -> None:
    _install_fake_upscale_runtime(monkeypatch)
    source = tmp_path / "source.png"
    output_root = tmp_path / "output"
    Image.new("RGB", (2, 2), color="red").save(source)

    image_upscale.main(inputs=[source], out_dir=output_root, cpu=True)

    output = capsys.readouterr().out
    assert "Model: Kim2091/UltraSharpV2/4x-test.safetensors" in output
    assert "Device: cpu | dtype: torch.float32 | scale: 4x" in output
    assert f"[1/1] Wrote: {output_root.resolve() / 'source.png'}" in output