tests/test_studio.py

from __future__ import annotations

from pathlib import Path

from adam.models import ExecutionPlan, PlanStep
from adam.studio import (
    PreviewEvaluation,
    StudioStore,
    TrainingRecipe,
    caption_path,
    exact_duplicate_groups,
    image_files,
)
from adam.training_assistant import estimate_plan


def test_dataset_review_and_recipe_store_round_trip(tmp_path: Path) -> None:
    dataset = tmp_path / "dataset"
    dataset.mkdir()
    first = dataset / "first.png"
    second = dataset / "second.png"
    first.write_bytes(b"same image bytes")
    second.write_bytes(b"same image bytes")
    caption_path(first).write_text("subject portrait\n", encoding="utf-8")

    assert image_files(dataset) == [first, second]
    assert exact_duplicate_groups([first, second]) == [[str(first), str(second)]]

    store = StudioStore(tmp_path)
    store.set_decision(str(dataset), str(first), "keep")
    recipe = store.add_recipe(
        TrainingRecipe("Portrait recipe", "lora", 80, preview_prompt="portrait")
    )
    store.add_evaluation(
        PreviewEvaluation("model-1", "checkpoint", "portrait", 42, 4)
    )
    assert store.toggle_best("model-1")

    restored = StudioStore(tmp_path)
    review = restored.review(str(dataset))
    assert review.decisions[str(first.resolve())] == "keep"
    assert restored.recipes[0].id == recipe.id
    assert restored.evaluations[0].seed == 42
    assert "model-1" in restored.best_models


def test_plan_estimate_is_labelled_and_bounded(tmp_path: Path) -> None:
    dataset = tmp_path / "images"
    dataset.mkdir()
    for index in range(3):
        (dataset / f"{index}.jpg").write_bytes(b"image")
    plan = ExecutionPlan(
        request="train",
        summary="Train",
        steps=[
            PlanStep(
                "lora_trainer",
                "Train LoRA",
                "Training",
                {"dataset_dir": str(dataset), "epochs": 10},
            )
        ],
    )

    estimates = estimate_plan(plan)

    assert len(estimates) == 3
    assert all(item.level == "estimate" for item in estimates)
    assert "30 image-epochs" in estimates[0].message
    assert "Rough duration" in estimates[1].message
    assert "VRAM" in estimates[2].message


def test_rejected_images_are_quarantined_and_recoverable(tmp_path: Path) -> None:
    dataset = tmp_path / "dataset"
    dataset.mkdir()
    image = dataset / "bad.png"
    image.write_bytes(b"bad image")
    caption_path(image).write_text("wrong subject", encoding="utf-8")
    store = StudioStore(tmp_path)
    store.set_decision(str(dataset), str(image), "reject")

    assert store.apply_rejections(str(dataset)) == 1
    assert not image.exists()
    assert not caption_path(image).exists()
    assert store.restore_rejections(str(dataset)) == 1
    assert image.is_file()
    assert caption_path(image).is_file()


def test_dataset_review_can_keep_all_then_reject_one(tmp_path: Path) -> None:
    dataset = tmp_path / "dataset"
    dataset.mkdir()
    images = [dataset / "one.png", dataset / "two.png", dataset / "three.png"]
    for image in images:
        image.write_bytes(b"image")
    store = StudioStore(tmp_path)

    assert store.set_all_decisions(str(dataset), images, "keep") == 3
    store.set_decision(str(dataset), str(images[1]), "reject")

    decisions = store.review(str(dataset)).decisions
    assert decisions[str(images[0].resolve())] == "keep"
    assert decisions[str(images[1].resolve())] == "reject"
    assert decisions[str(images[2].resolve())] == "keep"


def test_dataset_review_applies_mixed_decisions_in_one_operation(tmp_path: Path) -> None:
    dataset = tmp_path / "dataset"
    dataset.mkdir()
    first = dataset / "first.png"; first.write_bytes(b"one")
    second = dataset / "second.png"; second.write_bytes(b"two")
    store = StudioStore(tmp_path)

    changed = store.apply_decisions(
        str(dataset), {str(first): "keep", str(second): "reject"}
    )

    assert changed == 2
    assert store.review(str(dataset)).decisions[str(first.resolve())] == "keep"
    assert store.review(str(dataset)).decisions[str(second.resolve())] == "reject"