tests/test_cnn_reviewer.py

from pathlib import Path

import pytest

from adam.cnn_reviewer import _labeled_paths, reviewer_checkpoint, score_frames, train_reviewer


def test_reviewer_checkpoint_is_adam_owned_and_dataset_specific(tmp_path: Path) -> None:
    first = reviewer_checkpoint(tmp_path, tmp_path / "dataset one")
    second = reviewer_checkpoint(tmp_path, tmp_path / "dataset two")

    assert first.parent == tmp_path / "data" / "cnn_reviewers"
    assert first != second
    assert first.suffix == ".pt"


def test_reviewer_uses_only_readable_image_keep_reject_labels(tmp_path: Path) -> None:
    kept = tmp_path / "keep.png"; kept.write_bytes(b"not decoded here")
    rejected = tmp_path / "reject.jpg"; rejected.write_bytes(b"not decoded here")
    missing = tmp_path / "missing.png"

    paths, labels = _labeled_paths({str(kept): "keep", str(rejected): "reject", str(missing): "keep", "x": "unreviewed"})

    assert paths == [kept, rejected]
    assert labels == [1.0, 0.0]


def test_reviewer_trains_and_scores_without_touching_decisions(tmp_path: Path) -> None:
    pytest.importorskip("torch")
    from PIL import Image

    decisions = {}
    for index in range(8):
        path = tmp_path / f"keep_{index}.png"
        Image.new("RGB", (32, 32), (20, 220, 30)).save(path)
        decisions[str(path)] = "keep"
    for index in range(8):
        path = tmp_path / f"reject_{index}.png"
        Image.new("RGB", (32, 32), (220, 20, 30)).save(path)
        decisions[str(path)] = "reject"

    result = train_reviewer(tmp_path, tmp_path, decisions, epochs=1)
    scores = score_frames(result.checkpoint, tmp_path)

    assert Path(result.checkpoint).is_file()
    assert len(scores) == 16
    assert decisions[str(tmp_path / "keep_0.png")] == "keep"