tests/test_image_caption_progress.py

from __future__ import annotations

from pathlib import Path
from types import SimpleNamespace

from PIL import Image
import pytest

import kneiff.utils.image.caption.blip as caption_blip
import kneiff.utils.image.caption.server as caption_server
from kneiff.progress import ProgressUpdate
from kneiff.utils.image.caption.batch import (
    CaptionBatchItemResult,
    plan_caption_batch,
)
from kneiff.utils.image.caption.io import CaptionSidecarIO


def test_caption_batch_plan_excludes_existing_sidecars(tmp_path: Path) -> None:
    first = tmp_path / "first.png"
    second = tmp_path / "second.png"
    Image.new("RGB", (2, 2), color="red").save(first)
    Image.new("RGB", (2, 2), color="blue").save(second)
    first.with_suffix(".cap.txt").write_text("existing", encoding="utf-8")

    jobs = plan_caption_batch(tmp_path, CaptionSidecarIO())

    assert [job.image_path for job in jobs] == [second]


def test_server_caption_batch_reports_written_and_caught_failures(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    first = tmp_path / "first.png"
    second = tmp_path / "second.png"
    Image.new("RGB", (2, 2), color="red").save(first)
    Image.new("RGB", (2, 2), color="blue").save(second)
    sidecar_io = CaptionSidecarIO()
    jobs = plan_caption_batch(tmp_path, sidecar_io)
    calls = 0

    monkeypatch.setattr(
        caption_server,
        "load_openai_backend_config",
        lambda **kwargs: SimpleNamespace(model="test-model"),
    )
    monkeypatch.setattr(caption_server, "build_openai_client", lambda cfg: object())
    monkeypatch.setattr(
        caption_server,
        "image_to_data_url",
        lambda *args, **kwargs: "data:image/jpeg;base64,test",
    )

    def fake_caption_one(*args: object, **kwargs: object) -> str:
        nonlocal calls
        calls += 1
        if calls == 2:
            raise RuntimeError("server unavailable")
        return "A concise caption."

    monkeypatch.setattr(caption_server, "caption_one", fake_caption_one)
    updates: list[ProgressUpdate] = []
    items: list[CaptionBatchItemResult] = []
    instructions: list[str] = []

    result = caption_server.run_caption_batch_with_server(
        jobs,
        sidecar_io=sidecar_io,
        model_name="test-model",
        progress_callback=updates.append,
        instruction_callback=instructions.append,
        item_callback=items.append,
    )

    assert result.count("written") == 1
    assert result.count("failed") == 1
    assert first.with_suffix(".cap.txt").is_file()
    assert not second.with_suffix(".cap.txt").exists()
    assert len(instructions) == 2
    assert [item.status for item in items] == ["written", "failed"]
    assert sum(update.advance for update in updates) == 2


class _FakeCaptioner:
    def caption(self, image: Image.Image, instruction: str, gen: object) -> str:
        del image, instruction, gen
        return "A local caption."


def test_blip_caption_batch_advances_for_image_open_skip(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    invalid = tmp_path / "invalid.png"
    valid = tmp_path / "valid.png"
    invalid.write_bytes(b"not an image")
    Image.new("RGB", (2, 2), color="green").save(valid)
    sidecar_io = CaptionSidecarIO()
    jobs = plan_caption_batch(tmp_path, sidecar_io)
    monkeypatch.setattr(
        caption_blip,
        "build_captioner",
        lambda model_name, backend: _FakeCaptioner(),
    )
    updates: list[ProgressUpdate] = []

    result = caption_blip.run_caption_batch_with_blip(
        jobs,
        sidecar_io=sidecar_io,
        progress_callback=updates.append,
    )

    assert result.count("skipped") == 1
    assert result.count("written") == 1
    assert valid.with_suffix(".cap.txt").is_file()
    assert sum(update.advance for update in updates) == 2


def test_blip_inference_failure_propagates_with_incomplete_progress(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    source = tmp_path / "source.png"
    Image.new("RGB", (2, 2), color="green").save(source)
    sidecar_io = CaptionSidecarIO()
    jobs = plan_caption_batch(tmp_path, sidecar_io)

    class FailingCaptioner:
        def caption(self, image: Image.Image, instruction: str, gen: object) -> str:
            del image, instruction, gen
            raise RuntimeError("inference failed")

    monkeypatch.setattr(
        caption_blip,
        "build_captioner",
        lambda model_name, backend: FailingCaptioner(),
    )
    updates: list[ProgressUpdate] = []

    with pytest.raises(RuntimeError, match="inference failed"):
        caption_blip.run_caption_batch_with_blip(
            jobs,
            sidecar_io=sidecar_io,
            progress_callback=updates.append,
        )

    assert updates == []
    assert not source.with_suffix(".cap.txt").exists()