tests/test_cli_image_progress.py

from __future__ import annotations

from io import StringIO
from pathlib import Path

from PIL import Image
import pytest
from rich.console import Console
from typer.testing import CliRunner

import kneiff.cli.image as cli_image
import kneiff.utils.image.caption.server as caption_server
import kneiff.utils.image.tag_jtp3 as tag_jtp3
from kneiff.progress import ProgressUpdate
from kneiff.utils.image.caption.batch import (
    CaptionBatchItemResult,
    CaptionBatchResult,
)
from tests._cli_helpers import invoke_cli


pytestmark = pytest.mark.usefixtures("isolated_cli_project")


def _reject_progress(*args: object, **kwargs: object) -> None:
    del args, kwargs
    raise AssertionError("This command must not create a progress reporter.")


def test_concat_and_rename_do_not_create_progress(
    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)
    monkeypatch.setattr(cli_image, "CliProgress", _reject_progress)

    concat_result = invoke_cli(
        CliRunner(),
        ["img", "concat", str(first), str(second), "-o", str(tmp_path / "out.png")],
    )
    rename_result = invoke_cli(
        CliRunner(),
        ["img", "rename", str(first), "--base-name", "sample"],
    )

    assert concat_result.exit_code == 0, concat_result.output
    assert rename_result.exit_code == 0, rename_result.output


def test_tag_does_not_create_progress(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    source = tmp_path / "source.png"
    source.write_bytes(b"image")
    monkeypatch.setattr(cli_image, "CliProgress", _reject_progress)
    monkeypatch.setattr(tag_jtp3, "run_jtp3", lambda *args, **kwargs: None)

    result = invoke_cli(CliRunner(), ["img", "tag", str(source)])

    assert result.exit_code == 0, result.output


def test_single_image_caption_does_not_create_progress(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    source = tmp_path / "source.png"
    source.write_bytes(b"image")
    monkeypatch.setattr(cli_image, "CliProgress", _reject_progress)
    monkeypatch.setattr(
        caption_server,
        "caption_image_with_server",
        lambda **kwargs: "one caption",
    )

    result = invoke_cli(CliRunner(), ["img", "caption", str(source)])

    assert result.exit_code == 0, result.output
    assert result.output.strip() == "one caption"


def test_directory_caption_redirected_output_keeps_instruction_and_dry_run(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    source = tmp_path / "source.png"
    source.write_bytes(b"image")

    def fake_run(jobs, **kwargs: object) -> CaptionBatchResult:
        job = jobs[0]
        instruction_callback = kwargs["instruction_callback"]
        item_callback = kwargs["item_callback"]
        progress_callback = kwargs["progress_callback"]
        assert callable(instruction_callback)
        assert callable(item_callback)
        assert callable(progress_callback)
        instruction_callback("Instruction:\nDescribe the image.\n")
        item = CaptionBatchItemResult(
            job=job,
            status="planned",
            caption="one dry-run caption",
        )
        item_callback(item)
        progress_callback(ProgressUpdate("Caption 1/1 · planned", advance=1))
        return CaptionBatchResult(items=(item,))

    monkeypatch.setattr(caption_server, "run_caption_batch_with_server", fake_run)

    result = invoke_cli(
        CliRunner(),
        ["img", "caption", str(tmp_path), "--dry-run"],
    )

    assert result.exit_code == 0, result.output
    assert "Instruction:\nDescribe the image." in result.output
    assert "source.png -> one dry-run caption" in result.output
    assert "\x1b[" not in result.output


def test_directory_caption_terminal_hides_instruction_dump(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    source = tmp_path / "source.png"
    source.write_bytes(b"image")
    terminal_output = StringIO()
    console = Console(
        file=terminal_output,
        force_terminal=True,
        color_system=None,
        width=100,
    )
    monkeypatch.setattr(cli_image, "stderr_progress_console", lambda: console)

    def fake_run(jobs, **kwargs: object) -> CaptionBatchResult:
        assert kwargs["instruction_callback"] is None
        progress_callback = kwargs["progress_callback"]
        assert callable(progress_callback)
        item = CaptionBatchItemResult(job=jobs[0], status="written", caption="caption")
        progress_callback(ProgressUpdate("Caption 1/1", advance=1))
        return CaptionBatchResult(items=(item,))

    monkeypatch.setattr(caption_server, "run_caption_batch_with_server", fake_run)

    result = invoke_cli(CliRunner(), ["img", "caption", str(tmp_path)])

    assert result.exit_code == 0, result.output
    assert "Instruction:" not in result.output
    assert "Captioning complete" in terminal_output.getvalue()