tests/test_job_manager.py

from __future__ import annotations

import logging
from datetime import datetime, timedelta, timezone
from pathlib import Path

from adam.job_manager import JobManager, JobWorker
from adam.models import ExecutionPlan, Job, JobStatus, PlanStep
from adam.training_assistant import append_preflight_summary


def test_submission_reviews_training_before_queueing(tmp_path: Path, monkeypatch) -> None:
    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))
    starts = []
    monkeypatch.setattr(manager, "_start_next", lambda: starts.append(True))
    dataset = tmp_path / "planned_dataset"
    plan = ExecutionPlan(
        request="collect and train", summary="Collect and train.",
        steps=[
            PlanStep("dataset_collector", "Collect", "Collect", {
                "output_dir": str(dataset), "image_count": 2000,
            }),
            PlanStep("ddpm_trainer", "Train", "Train", {
                "dataset_dir": str(dataset), "epochs": 600, "batch_size": 1,
            }),
        ],
    )

    job = manager.submit(plan)

    assert job.status == JobStatus.AWAITING_CONFIRMATION
    assert plan.orion_review["level"] == "warning"
    assert "Pre-flight:" in plan.summary
    assert starts == []
    assert manager._queue == []
    restored = JobManager(tmp_path, None, logging.getLogger("test.jobs"))
    assert restored.jobs[0].plan.orion_review == plan.orion_review
    assert restored.jobs[0].status == JobStatus.AWAITING_CONFIRMATION


def test_submission_preserves_an_already_reviewed_plan(tmp_path: Path) -> None:
    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))
    plan = ExecutionPlan(
        request="train", summary="Train.", requires_confirmation=True,
        steps=[PlanStep("ddpm_trainer", "Train", "Train", {"epochs": 10})],
    )
    append_preflight_summary(plan, {})
    summary = plan.summary
    arguments = dict(plan.steps[0].arguments)

    job = manager.submit(plan)

    assert job.plan.summary == summary
    assert job.plan.steps[0].arguments == arguments
    assert summary.count("ORION —") == 1
    assert summary.count("Pre-flight:") == 1


def _job(index: int, status: JobStatus = JobStatus.FINISHED) -> Job:
    return Job(
        id=f"JOB{index:04d}",
        plan=ExecutionPlan(
            request="test",
            summary="Test job",
            steps=[],
            project_name=f"Project {index}",
        ),
        status=status,
    )


def test_job_history_is_not_silently_truncated(tmp_path: Path) -> None:
    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    manager.jobs = [_job(index) for index in range(125)]
    manager._save()

    restored = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]

    assert len(restored.jobs) == 125


def test_jobs_history_rows_remain_visible_while_loading(tmp_path: Path, monkeypatch) -> None:
    """A long history must not look blank while its rows are streamed into the table."""
    monkeypatch.setenv("QT_QPA_PLATFORM", "offscreen")
    from PySide6.QtWidgets import QApplication
    from PySide6.QtTest import QTest
    from adam.ui.main_window import JobsPage

    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    manager.jobs = [_job(index) for index in range(125)]
    app = QApplication.instance() or QApplication([])
    page = JobsPage(manager)

    assert page.table.updatesEnabled()
    QTest.qWait(200)
    assert page.table.item(0, 0).text() == "JOB0000"
    page.close()
    page.deleteLater()
    app.processEvents()


def test_jobs_page_can_approve_a_waiting_plan(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setenv("QT_QPA_PLATFORM", "offscreen")
    from PySide6.QtWidgets import QApplication
    from adam.ui.main_window import JobsPage

    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    waiting = _job(1, JobStatus.AWAITING_CONFIRMATION)
    manager.jobs = [waiting]
    approved: list[str] = []
    monkeypatch.setattr(manager, "confirm", approved.append)
    app = QApplication.instance() or QApplication([])
    page = JobsPage(manager)
    page.show()
    page.selected_job_id = waiting.id
    page._show_job(waiting)

    assert page.approve_button.isVisible()
    assert page.approve_button.isEnabled()
    page.approve_button.click()
    assert approved == [waiting.id]
    page.close()
    page.deleteLater()
    app.processEvents()


def test_history_cleanup_only_removes_completed_and_failed(tmp_path: Path) -> None:
    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    manager.jobs = [
        _job(1, JobStatus.FINISHED),
        _job(2, JobStatus.FAILED),
        _job(3, JobStatus.CANCELLED),
        _job(4, JobStatus.INTERRUPTED),
    ]

    assert manager.remove_completed_or_failed() == 2
    assert [job.status for job in manager.jobs] == [
        JobStatus.CANCELLED,
        JobStatus.INTERRUPTED,
    ]


def test_end_task_acknowledges_an_interrupted_job(tmp_path: Path) -> None:
    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    job = _job(1, JobStatus.INTERRUPTED)
    job.logs.append("[startup] Previous session ended before this job. Review it before retrying.")
    manager.jobs = [job]

    assert manager.end_task(job.id) is True
    assert job.status == JobStatus.CANCELLED
    assert not any("Previous session ended" in line for line in job.logs)

    restored = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    assert restored.jobs[0].status == JobStatus.CANCELLED


def test_approved_future_job_stays_scheduled_and_survives_restart(tmp_path: Path) -> None:
    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    plan = ExecutionPlan(
        request="train later", summary="Scheduled training", steps=[],
        requires_confirmation=True,
    )
    start = (datetime.now(timezone.utc) + timedelta(hours=2)).isoformat()

    job = manager.submit(plan, scheduled_for=start)
    assert job.status == JobStatus.AWAITING_CONFIRMATION
    manager.confirm(job.id)
    assert job.status == JobStatus.SCHEDULED

    restored = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    assert restored.jobs[0].status == JobStatus.SCHEDULED
    assert restored.jobs[0].scheduled_for == start


def test_due_schedule_queues_behind_an_active_job(tmp_path: Path) -> None:
    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    job = _job(2, JobStatus.SCHEDULED)
    job.scheduled_for = (datetime.now(timezone.utc) - timedelta(minutes=1)).isoformat()
    manager.jobs = [job]

    class BusyWorker:
        @staticmethod
        def isRunning() -> bool:
            return True

    manager._worker = BusyWorker()  # type: ignore[assignment]
    manager._release_due_scheduled()

    assert job.status == JobStatus.QUEUED
    assert manager._queue == [job.id]


def test_restart_preserves_queued_jobs_and_interrupts_only_active_work(tmp_path: Path) -> None:
    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    queued = _job(1, JobStatus.QUEUED)
    running = _job(2, JobStatus.RUNNING)
    manager.jobs = [queued, running]
    manager._save()

    restored = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]

    assert restored.jobs[0].status == JobStatus.QUEUED
    assert restored.jobs[1].status == JobStatus.INTERRUPTED
    assert restored._queue == [queued.id]


def test_worker_coalesces_rapid_progress_events() -> None:
    plan = ExecutionPlan(
        request="train",
        summary="Training",
        steps=[PlanStep("ddpm_trainer", "Train", "Run training")],
    )
    job = Job(id="FAST0001", plan=plan, status=JobStatus.RUNNING)

    class NoisyExecutor:
        def execute(self, _tool_id, _arguments, **kwargs):
            for percent in range(1, 101):
                kwargs["progress_callback"](percent, "same training burst")
            return {}

    worker = JobWorker(job, NoisyExecutor())  # type: ignore[arg-type]
    events = []
    worker.event.connect(events.append)

    worker.run()

    progress_events = [event for event in events if event.get("type") == "progress"]
    assert 1 <= len(progress_events) <= 2
    assert progress_events[-1]["overall"] == 100


def test_step_eta_uses_measured_progress_cadence() -> None:
    samples: list[dict[str, object]] = []

    first = JobWorker._estimate_step_eta(
        {"current_step": 10, "total_steps": 110, "unit": "step"},
        samples,
        100.0,
    )
    second = JobWorker._estimate_step_eta(
        {"current_step": 20, "total_steps": 110, "unit": "step"},
        samples,
        120.0,
    )

    assert "eta_seconds" not in first
    assert second["eta_seconds"] == 180
    assert second["progress_current"] == 20
    assert second["progress_total"] == 110
    assert second["progress_unit"] == "step"
    assert second["progress_rate"] == 0.5
    assert second["estimated_completion_at"]


def test_step_eta_resets_when_progress_counter_changes() -> None:
    samples: list[dict[str, object]] = []
    JobWorker._estimate_step_eta(
        {"current_step": 819, "total_steps": 825, "unit": "step"}, samples, 100.0
    )
    reset = JobWorker._estimate_step_eta(
        {"current_step": 1, "total_steps": 160, "unit": "step"}, samples, 120.0
    )
    measured = JobWorker._estimate_step_eta(
        {"current_step": 2, "total_steps": 160, "unit": "step"}, samples, 141.5
    )

    assert "eta_seconds" not in reset
    assert measured["eta_seconds"] == 3397
    assert measured["progress_current"] == 2
    assert measured["progress_total"] == 160


def test_active_ddpm_adjustment_is_queued_on_worker(tmp_path: Path, monkeypatch) -> None:
    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    job = Job(
        plan=ExecutionPlan(
            request="train", summary="Train",
            steps=[PlanStep("ddpm_trainer", "Train", "Train", {
                "batch_size": 8, "gradient_accumulation_steps": 1,
                "training_intensity": 100, "epochs": 20,
            })],
        ),
        status=JobStatus.RUNNING,
        current_step=0,
    )

    class Worker:
        updates = None

        def request_adjustment(self, updates):
            self.updates = updates

    worker = Worker()
    manager.jobs = [job]
    manager._active_job = job
    manager._worker = worker  # type: ignore[assignment]
    monkeypatch.setattr(manager, "_save", lambda: None)

    manager.request_training_adjustment(job.id, {
        "batch_size": 4,
        "gradient_accumulation_steps": 2,
        "training_intensity": 75,
    })

    assert worker.updates == {
        "batch_size": 4,
        "gradient_accumulation_steps": 2,
        "training_intensity": 75,
    }
    assert job.plan.steps[0].arguments["batch_size"] == 8


def test_vram_retry_uses_old_batch_to_calculate_completed_epochs(tmp_path: Path, monkeypatch) -> None:
    dataset = tmp_path / "dataset"
    output = tmp_path / "output"
    checkpoint = output / "checkpoint-20"
    (checkpoint / "unet").mkdir(parents=True)
    dataset.mkdir()
    for index in range(40):
        (dataset / f"{index}.png").write_bytes(b"image")
    (checkpoint / "unet" / "diffusion_pytorch_model.safetensors").write_bytes(b"weights")
    (checkpoint / "optimizer.bin").write_bytes(b"optimizer")
    (checkpoint / "scheduler.bin").write_bytes(b"scheduler")
    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    failed = Job(
        plan=ExecutionPlan(
            request="train", summary="Train",
            steps=[PlanStep("ddpm_trainer", "Train", "Train", {
                "dataset_dir": str(dataset), "output_dir": str(output),
                "epochs": 10, "batch_size": 4, "gradient_accumulation_steps": 1,
            })],
        ),
        status=JobStatus.FAILED,
        current_step=0,
        error="CUDA out of memory",
    )
    manager.jobs = [failed]
    monkeypatch.setattr(manager, "_start_next", lambda: None)

    retry = manager.safer_vram_retry(failed.id)
    arguments = retry.plan.steps[0].arguments

    assert arguments["batch_size"] == 2
    assert arguments["gradient_accumulation_steps"] == 2
    assert arguments["completed_epochs"] == 2
    assert arguments["epochs"] == 8
    assert arguments["resume_from"] == str(checkpoint)


def test_adjustment_ready_requeues_same_job_from_checkpoint(tmp_path: Path, monkeypatch) -> None:
    manager = JobManager(tmp_path, None, logging.getLogger("test.jobs"))  # type: ignore[arg-type]
    job = Job(
        plan=ExecutionPlan(
            request="train", summary="Train",
            steps=[PlanStep("ddpm_trainer", "Train", "Train", {
                "epochs": 20, "batch_size": 8, "training_intensity": 100,
            })],
        ),
        status=JobStatus.RUNNING,
        current_step=0,
    )
    manager.jobs = [job]
    manager._active_job = job
    monkeypatch.setattr(manager, "_save", lambda: None)

    manager._handle_event({
        "type": "adjustment_ready",
        "checkpoint": str(tmp_path / "checkpoint-40"),
        "completed_epochs": 4,
        "updates": {"batch_size": 4, "training_intensity": 75},
    })

    arguments = job.plan.steps[0].arguments
    assert job.status == JobStatus.QUEUED
    assert manager._queue == [job.id]
    assert arguments["epochs"] == 16
    assert arguments["completed_epochs"] == 4
    assert arguments["batch_size"] == 4
    assert arguments["training_intensity"] == 75