llm-connect/tests/test_image_selection.py

"""Decision tree and hall portrait planning."""

import json
from pathlib import Path

import pytest

from llm_connect.cli import main
from llm_connect.exceptions import LLMConfigurationError
from llm_connect.hall_portraits import plan_missing_portraits, reveal_portrait
from llm_connect.image_selection import classify_image_branch, selection_for_prompt

HALL_PROMPT = (
    "Constellation dialect. Square, dark indigo. A gold-wire helix. "
    "Fine technical-illustration linework, no logos, no readable text."
)


def test_hall_house_prompt_is_general_efficient() -> None:
    selection = selection_for_prompt(HALL_PROMPT)
    assert classify_image_branch(HALL_PROMPT) == "general"
    assert selection.tier == "efficient"
    assert selection.model == "google/gemini-nano-banana-2.1"
    assert selection.quote_usd == 0.04
    assert selection.max_cost_usd == 0.08


def test_negated_text_does_not_become_text_heavy() -> None:
    assert classify_image_branch("no logos, no readable text, pale gold") == "general"


def test_abstract_marks_stay_on_the_general_branch() -> None:
    prompt = (
        "a ledger glows with the words rendered as abstract marks, not readable text. "
        "no logos, no lettering"
    )
    assert classify_image_branch(prompt) == "general"


def test_requested_words_take_the_text_branch() -> None:
    selection = selection_for_prompt("A poster whose headline is the words HELIX")
    assert selection.branch == "text_heavy"
    assert selection.model == "openai/gpt-image-2.5-sunburst"


def test_vector_and_self_host_follow_the_tree() -> None:
    assert selection_for_prompt("Export as a vector svg icon").model.startswith("recraft/")
    try:
        selection_for_prompt("self-host the open weights on our own hardware")
    except Exception as exc:
        assert "OpenRouter" in str(exc)
    else:
        raise AssertionError("self-host branch should refuse")


def _entry(visual: str, comment: str | None, entry_id: str = "hall-worker-ada-1") -> str:
    commented = f"<!-- ![{comment}](../visuals/ada.jpg) -->\n" if comment else ""
    return f"""---
id: {entry_id}
status: draft
---

# Ada — a seat

## Visual prompt

> {visual}

{commented}
## Handoff

Next.
"""


def test_plan_skips_a_portrait_that_is_already_on_disk(tmp_path: Path) -> None:
    entries = tmp_path / "entries"
    visuals = tmp_path / "visuals"
    entries.mkdir()
    visuals.mkdir()
    (entries / "one.md").write_text(_entry(HALL_PROMPT, "Ada"), encoding="utf-8")
    (entries / "two.md").write_text(_entry(HALL_PROMPT, None, "hall-worker-bea-2"), encoding="utf-8")
    (visuals / "bea-file.jpg").write_bytes(b"x")
    # two.md has no comment, so its derived file is the id. Put that file on disk.
    (visuals / "hall-worker-bea-2.jpg").write_bytes(b"x")

    jobs = plan_missing_portraits(tmp_path)
    assert [job.entry_id for job in jobs] == ["hall-worker-ada-1"]
    assert jobs[0].output_path == visuals / "ada.jpg"
    assert jobs[0].relative_path == "../visuals/ada.jpg"
    assert jobs[0].alt == "Ada"
    assert "no readable text" in jobs[0].prompt


def test_reveal_portrait_uncomments_the_intended_line() -> None:
    text = _entry(HALL_PROMPT, "Ada")
    updated = reveal_portrait(text, "../visuals/ada.jpg", "Ada")
    assert "<!--" not in updated
    assert "![Ada](../visuals/ada.jpg)" in updated
    assert "status: draft" in updated


def test_plan_reads_a_prompt_that_is_not_a_blockquote(tmp_path: Path) -> None:
    entries = tmp_path / "entries"
    visuals = tmp_path / "visuals"
    entries.mkdir()
    visuals.mkdir()
    body = """---
id: hall-worker-cy-3
status: draft
---

# Cy — a seat

## Visual prompt

**Dialect:** Constellation. Square, no logos, no readable text.

A gold-wire orrery suspended in dark indigo.

## Handoff

Next.
"""
    (entries / "three.md").write_text(body, encoding="utf-8")
    jobs = plan_missing_portraits(tmp_path)
    assert len(jobs) == 1
    assert "orrery" in jobs[0].prompt
    assert jobs[0].output_path.name == "hall-worker-cy-3.jpg"


def test_a_same_named_file_outside_entries_is_not_the_seat(tmp_path: Path, monkeypatch) -> None:
    hall = tmp_path / "hall"
    entries = hall / "entries"
    visuals = hall / "visuals"
    entries.mkdir(parents=True)
    visuals.mkdir()
    (tmp_path / "one.md").write_text("not a hall entry\n", encoding="utf-8")
    (entries / "one.md").write_text(_entry(HALL_PROMPT, "Ada"), encoding="utf-8")
    monkeypatch.chdir(tmp_path)
    jobs = plan_missing_portraits(hall, entry="one.md")
    assert len(jobs) == 1
    assert jobs[0].entry_path == (entries / "one.md").resolve()
    with pytest.raises(LLMConfigurationError):
        plan_missing_portraits(hall, entry=tmp_path / "one.md")


def test_one_entry_is_selected_for_publication(tmp_path: Path, capsys) -> None:
    entries = tmp_path / "entries"
    visuals = tmp_path / "visuals"
    entries.mkdir()
    visuals.mkdir()
    (entries / "one.md").write_text(_entry(HALL_PROMPT, "Ada"), encoding="utf-8")
    (entries / "two.md").write_text(_entry(HALL_PROMPT, None, "hall-worker-bea-2"), encoding="utf-8")
    try:
        main(["image", "render-hall", "--hall", str(tmp_path), "--dry-run"])
    except SystemExit:
        pass
    else:
        raise AssertionError("render-hall without --entry should fail")
    code = main(
        [
            "image",
            "render-hall",
            "--hall",
            str(tmp_path),
            "--entry",
            "one.md",
            "--dry-run",
            "--json",
        ]
    )
    assert code == 0
    payload = json.loads(capsys.readouterr().out)
    assert payload["count"] == 1
    assert payload["jobs"][0]["entry_id"] == "hall-worker-ada-1"
    assert not (visuals / "ada.jpg").exists()


def test_cli_plan_is_dry_and_names_the_efficient_model(tmp_path: Path, capsys) -> None:
    entries = tmp_path / "entries"
    visuals = tmp_path / "visuals"
    entries.mkdir()
    visuals.mkdir()
    (entries / "one.md").write_text(_entry(HALL_PROMPT, "Ada"), encoding="utf-8")
    code = main(["image", "plan", "--hall", str(tmp_path), "--json"])
    assert code == 0
    payload = json.loads(capsys.readouterr().out)
    assert payload["count"] == 1
    assert payload["jobs"][0]["model"] == "google/gemini-nano-banana-2.1"
    assert payload["max_batch_cost_usd"] == 0.08
    assert not (visuals / "ada.jpg").exists()