tests/test_render_visual_tables.py

from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path

import pytest
from PIL import Image
from render_visual_tables import CODECS, load_prompt_ids, load_rows, render, render_contact_sheet


def test_load_prompt_ids_accepts_manifest_object(tmp_path: Path) -> None:
    path = tmp_path / "prompts.json"
    path.write_text(json.dumps({"prompt_ids": ["p2", "p1"]}), encoding="utf-8")

    assert load_prompt_ids(path) == ["p2", "p1"]


def test_load_rows_fails_closed_when_a_success_has_no_sample(tmp_path: Path) -> None:
    run = tmp_path / "run"
    (run / "results").mkdir(parents=True)
    row = {"row_type": "generation", "condition": "png_baseline", "prompt_id": "p1", "success": True}
    (run / "results/raw_rows.jsonl").write_text(json.dumps(row) + "\n", encoding="utf-8")

    with pytest.raises(ValueError, match="sample"):
        load_rows(run, ["p1"])


def test_load_rows_rejects_raw_rows_symlink(tmp_path: Path) -> None:
    run = tmp_path / "run"
    (run / "results").mkdir(parents=True)
    target = tmp_path / "raw_rows.jsonl"
    target.write_text("", encoding="utf-8")
    (run / "results/raw_rows.jsonl").symlink_to(target)

    with pytest.raises(ValueError, match="symlink"):
        load_rows(run, ["p1"])


def test_load_rows_rejects_raw_rows_path_escape(tmp_path: Path) -> None:
    run = tmp_path / "run"
    (run / "results").mkdir(parents=True)
    (run / "results/raw_rows.jsonl").symlink_to(tmp_path / "outside.jsonl")
    (tmp_path / "outside.jsonl").write_text("", encoding="utf-8")

    with pytest.raises(ValueError, match="escapes run|symlink"):
        load_rows(run, ["p1"])


def test_render_rejects_sample_replaced_with_symlink_after_validation(tmp_path: Path) -> None:
    run = tmp_path / "run"
    sample = run / "samples/sample.png"
    sample.parent.mkdir(parents=True)
    Image.new("RGB", (4, 4), "blue").save(sample)
    sample_hash = hashlib.sha256(sample.read_bytes()).hexdigest()
    conditions = {"png_baseline", "webp_lossless"} | {
        condition for codec in CODECS.values() for condition in codec[1:]
    }
    rows = [
        {
            "row_type": "generation",
            "condition": condition,
            "prompt_id": "p1",
            "prompt": "test prompt",
            "category": "test",
            "success": True,
            "sample_path": "samples/sample.png",
            "sample_hash": sample_hash,
            "clip_score": 1.0,
        }
        for condition in sorted(conditions)
    ]
    results = run / "results"
    results.mkdir()
    (results / "raw_rows.jsonl").write_text(
        "".join(json.dumps(row) + "\n" for row in rows),
        encoding="utf-8",
    )
    loaded = load_rows(run, ["p1"])
    outside = tmp_path / "outside.png"
    Image.new("RGB", (4, 4), "red").save(outside)
    sample.unlink()
    sample.symlink_to(outside)

    with pytest.raises(ValueError, match="symlink"):
        render(
            "jpeg_q100",
            CODECS["jpeg_q100"],
            run.resolve(),
            [loaded[("png_baseline", "p1")]],
            loaded,
            tmp_path / "output",
        )


def test_render_rejects_regular_sample_changed_after_validation(tmp_path: Path) -> None:
    run = tmp_path / "run"
    sample = run / "samples/sample.png"
    sample.parent.mkdir(parents=True)
    Image.new("RGB", (4, 4), "blue").save(sample)
    sample_hash = hashlib.sha256(sample.read_bytes()).hexdigest()
    conditions = {"png_baseline", "webp_lossless"} | {
        condition for codec in CODECS.values() for condition in codec[1:]
    }
    rows = [
        {
            "row_type": "generation",
            "condition": condition,
            "prompt_id": "p1",
            "prompt": "test prompt",
            "category": "test",
            "success": True,
            "sample_path": "samples/sample.png",
            "sample_hash": sample_hash,
            "clip_score": 1.0,
        }
        for condition in sorted(conditions)
    ]
    results = run / "results"
    results.mkdir()
    (results / "raw_rows.jsonl").write_text(
        "".join(json.dumps(row) + "\n" for row in rows),
        encoding="utf-8",
    )
    loaded = load_rows(run, ["p1"])
    Image.new("RGB", (4, 4), "red").save(sample)

    with pytest.raises(ValueError, match="hash mismatch"):
        render(
            "jpeg_q100",
            CODECS["jpeg_q100"],
            run.resolve(),
            [loaded[("png_baseline", "p1")]],
            loaded,
            tmp_path / "output",
        )


@pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="FIFO is unavailable on this platform")
def test_render_rejects_fifo_replacement_without_blocking(tmp_path: Path) -> None:
    run = tmp_path / "run"
    sample = run / "samples/sample.png"
    sample.parent.mkdir(parents=True)
    Image.new("RGB", (4, 4), "blue").save(sample)
    sample_hash = hashlib.sha256(sample.read_bytes()).hexdigest()
    conditions = {"png_baseline", "webp_lossless"} | {
        condition for codec in CODECS.values() for condition in codec[1:]
    }
    rows = [
        {
            "row_type": "generation",
            "condition": condition,
            "prompt_id": "p1",
            "prompt": "test prompt",
            "category": "test",
            "success": True,
            "sample_path": "samples/sample.png",
            "sample_hash": sample_hash,
            "clip_score": 1.0,
        }
        for condition in sorted(conditions)
    ]
    results = run / "results"
    results.mkdir()
    (results / "raw_rows.jsonl").write_text(
        "".join(json.dumps(row) + "\n" for row in rows),
        encoding="utf-8",
    )
    loaded = load_rows(run, ["p1"])
    sample.unlink()
    os.mkfifo(sample)

    with pytest.raises(ValueError, match="not a regular file"):
        render(
            "jpeg_q100",
            CODECS["jpeg_q100"],
            run.resolve(),
            [loaded[("png_baseline", "p1")]],
            loaded,
            tmp_path / "output",
        )


def test_render_contact_sheet_stacks_tables_at_full_width(tmp_path: Path) -> None:
    first = tmp_path / "first.png"
    second = tmp_path / "second.png"
    Image.new("RGB", (100, 50), "red").save(first)
    Image.new("RGB", (80, 60), "blue").save(second)

    output = render_contact_sheet([first, second], tmp_path / "output", gap=20)

    with Image.open(output) as image:
        assert image.size == (100, 130)
        assert image.getpixel((0, 0)) == (255, 0, 0)
        assert image.getpixel((10, 60)) == (13, 20, 27)
        assert image.getpixel((10, 70)) == (0, 0, 255)
        assert image.getpixel((10, 129)) == (0, 0, 255)