tests/test_agents.py

from __future__ import annotations

from pathlib import Path
import logging
from types import SimpleNamespace

from PIL import Image

from adam.atlas import AtlasSupervisor
from adam.job_manager import JobManager
from adam.models import ExecutionPlan, Job, JobStatus, PlanStep, SystemSnapshot
from adam.nova import evaluate_job_output
from adam.orion import apply_orion_review, recommend_training_settings
from adam.ui.main_window import CommandCenterPage


def _training_job(tmp_path: Path) -> Job:
    return Job(
        plan=ExecutionPlan(
            request="train",
            summary="Train.",
            steps=[PlanStep("ddpm_trainer", "Train DDPM", "Train", {
                "dataset_dir": str(tmp_path), "epochs": 10, "batch_size": 2,
                "resolution": 128,
            })],
        ),
        status=JobStatus.RUNNING,
        current_step=0,
        progress=10,
    )


def test_orion_flags_small_dataset_preset_applied_to_large_dataset(tmp_path: Path) -> None:
    dataset = tmp_path / "dataset"
    dataset.mkdir()
    for index in range(1_000):
        (dataset / f"{index}.png").touch()
    plan = ExecutionPlan(
        request="train",
        summary="Train a model.",
        steps=[PlanStep("ddpm_trainer", "Train DDPM", "Train", {
            "dataset_dir": str(dataset), "epochs": 600, "batch_size": 1,
            "resolution": 128,
        })],
    )

    report = apply_orion_review(plan)

    assert report["level"] == "warning"
    assert report["settings_changed"] is False
    assert plan.requires_confirmation is True
    assert "ORION" in plan.summary
    assert "600,000 image exposures" in plan.summary


def test_orion_uses_planned_collection_size_before_dataset_exists(tmp_path: Path) -> None:
    dataset = tmp_path / "future-dataset"
    plan = ExecutionPlan(
        request="collect and train",
        summary="Collect and train.",
        steps=[
            PlanStep("dataset_collector", "Collect", "Collect", {
                "output_dir": str(dataset), "image_count": 2_000,
            }),
            PlanStep("ddpm_trainer", "Train", "Train", {
                "dataset_dir": str(dataset), "epochs": 600, "batch_size": 1,
                "resolution": 128,
            }),
        ],
    )

    report = apply_orion_review(plan)

    assert report["level"] == "warning"
    assert "1,200,000 image exposures" in plan.summary


def test_orion_warning_blocks_trusted_automation() -> None:
    class Config:
        def get(self, key: str, default=None):
            return key == "trusted_dataset_ddpm_automation"

    page = SimpleNamespace(config=Config())
    plan = ExecutionPlan(
        request="train", summary="Train.", requires_confirmation=True,
        steps=[PlanStep("dataset_collector", "Collect", "Collect"), PlanStep("ddpm_trainer", "Train", "Train")],
        orion_review={"level": "warning"},
    )

    assert CommandCenterPage._can_trusted_start(page, plan) is False


def test_orion_recipe_scales_epochs_and_batch_to_dataset_size_and_resolution() -> None:
    recipe = recommend_training_settings("ddpm", 1_000, 128, vram_gb=12)

    assert recipe["epochs"] == 180
    assert recipe["settings"]["batch_size"] == 12
    assert recipe["settings"]["preview_every"] == 18
    assert "180,000 image exposures" in recipe["summary"]


def test_orion_recipe_becomes_more_conservative_at_high_resolution() -> None:
    recipe = recommend_training_settings("flow", 300, 512, vram_gb=6)

    assert recipe["settings"]["batch_size"] == 1
    assert recipe["settings"]["gradient_checkpointing"] is True
    assert recipe["settings"]["workers"] == recipe["settings"]["dataloader_num_workers"]


def test_atlas_pauses_for_a_new_non_finite_loss_message(tmp_path: Path) -> None:
    job = _training_job(tmp_path)
    job.logs.append("loss = NaN")
    atlas = AtlasSupervisor()

    decision = atlas.observe(job, SystemSnapshot())

    assert decision.severity == "critical"
    assert decision.action == "pause"


def test_atlas_requires_repeated_critical_temperature_samples(tmp_path: Path) -> None:
    job = _training_job(tmp_path)
    atlas = AtlasSupervisor()
    hot = SystemSnapshot(gpu_temperature=91)

    assert atlas.observe(job, hot).action == "none"
    assert atlas.observe(job, hot).action == "none"
    assert atlas.observe(job, hot).action == "pause"


def test_job_manager_applies_an_atlas_critical_pause(tmp_path: Path) -> None:
    class Worker:
        paused = False

        def pause(self) -> None:
            self.paused = True

    manager = JobManager(tmp_path, None, logging.getLogger("test.atlas"))  # type: ignore[arg-type]
    job = _training_job(tmp_path)
    job.logs.append("loss: inf")
    worker = Worker()
    manager.jobs = [job]
    manager._active_job = job
    manager._worker = worker  # type: ignore[assignment]

    manager.supervise(SystemSnapshot())

    assert worker.paused is True
    assert job.status == JobStatus.PAUSED
    assert job.atlas_report["severity"] == "critical"


def test_nova_reports_duplicate_sample_collapse(tmp_path: Path) -> None:
    output = tmp_path / "output"
    output.mkdir()
    for index in range(4):
        Image.new("RGB", (32, 32), (80, 90, 100)).save(output / f"sample_{index}.png")
    job = _training_job(tmp_path)
    job.output_folder = str(output)

    report = evaluate_job_output(job)

    assert report["status"] == "NEEDS REVIEW"
    assert report["sample_count"] == 4
    assert report["duplicate_count"] == 3


def test_nova_requests_samples_when_training_has_no_generations(tmp_path: Path) -> None:
    job = _training_job(tmp_path)

    report = evaluate_job_output(job)

    assert report["status"] == "NEEDS SAMPLES"
    assert report["sample_count"] == 0