atlas-iac/testing/tests/test_suite_deadlines.py

"""Explicit long budgets share one deadline across routing, passes and progress."""
import json
import threading
from pathlib import Path

import pytest

import suite_backends
import suite_jobs
import suite_multipass
import suite_recovery
from suite_contract import MAX_JOB_SECONDS, Problem, TIMEOUT, preflight, validate_request
from test_suite_multipass import install_backend, natural, request


@pytest.mark.parametrize("seconds", [1801, 2700, 3600, 7200])
def test_explicit_long_budget_matches_published_schema(seconds):
    """Admission accepts opt-in durations; omitted limits retain their old default."""
    source = request(7)
    assert source["execution"]["max_seconds"] == TIMEOUT == 1800
    source["execution"]["max_seconds"] = seconds
    source = validate_request(source, ["claude"])
    assert preflight(source)["provider"] == "claude"
    schema = json.loads((Path(__file__).resolve().parents[2] /
                         "docs/contracts/suite_planning_request.schema.json").read_text())
    limit = schema["properties"]["execution"]["properties"]["max_seconds"]
    assert limit["maximum"] == MAX_JOB_SECONDS == 7200 and limit["default"] == TIMEOUT


@pytest.mark.parametrize("seconds", [7201, 7200.1, True, "7200", 9])
def test_invalid_budget_fails_before_inference(seconds):
    source = request(7)
    source["execution"]["max_seconds"] = seconds
    with pytest.raises(Problem, match="invalid_timeout"):
        validate_request(source, ["claude"])


def test_job_passes_1800_seconds_with_one_absolute_deadline(tmp_path, monkeypatch):
    """Advance synthetic time through a complete five-pass job beyond the old cap."""
    source = request(7)
    source["execution"]["max_seconds"] = 3600
    family = natural(source, [[c["alias"] for c in source["cases"]]])
    install_backend(monkeypatch, source, family)
    original = suite_backends.claude_generate
    clock, allocations, deadlines = [100.0], [], []
    monkeypatch.setattr(suite_jobs.time, "monotonic", lambda: clock[0])

    def routing(_):
        clock[0] += 10

    def backend(value, cancel, **kwargs):
        allocations.append(value["execution"]["max_seconds"])
        deadlines.append(kwargs["job_deadline"])
        result = original(value, cancel, **kwargs)
        clock[0] += 500
        return result

    monkeypatch.setattr(suite_backends, "switchyard_decision", routing)
    monkeypatch.setattr(suite_backends, "claude_generate", backend)
    jobs = suite_jobs.Jobs(tmp_path / "jobs.sqlite")
    selected = preflight(source)
    job, _ = jobs.submit("owner", "long-budget", source, selected, "192.168.22.8", launch=False)
    jobs.run(job["job_id"], "owner", source, selected, "192.168.22.8")
    result = jobs.get(job["job_id"], "owner")
    assert result["status"] == "completed" and result["wall_seconds"] == 2510
    assert allocations == [1795, 1743.75, 1692.5, 1641.25, 1590]
    assert deadlines == [3700] * 5
    assert result["execution_progress"]["job_remaining_seconds"] == 1090
    assert result["execution_progress"]["job_elapsed_seconds"] == 2510
    assert result["execution"]["max_seconds"] == 3600


@pytest.mark.parametrize("provider_error", ["backend_timeout_or_unavailable", "timeout"])
def test_distinct_backend_timeout_is_not_reclassified(monkeypatch, provider_error):
    """Only the shared-deadline watchdog identifies job time exhaustion."""
    source = request(7)
    def fail(*_, **__):
        raise Problem(provider_error, 504, failure_stage="subprocess")
    monkeypatch.setattr(suite_backends, "claude_generate", fail)
    monkeypatch.setattr(suite_multipass, "FALLBACK_ERRORS", set())
    monkeypatch.setattr(suite_recovery, "MAX_ATTEMPTS", 1)
    with pytest.raises(Problem, match="^(pass_provider_transient_failure|pass_timeout)$"):
        suite_multipass.generate(source, preflight(source), threading.Event(), "192.168.22.8")


def test_routing_consumes_job_budget_and_stops_before_provider(tmp_path, monkeypatch):
    source = request(7)
    source["execution"]["max_seconds"] = 10
    clock = [100.0]
    monkeypatch.setattr(suite_jobs.time, "monotonic", lambda: clock[0])
    monkeypatch.setattr(suite_backends, "switchyard_decision", lambda _: clock.__setitem__(0, 111))
    monkeypatch.setattr(suite_backends, "claude_generate", lambda *a, **k: pytest.fail("provider launched"))
    jobs = suite_jobs.Jobs(tmp_path / "jobs.sqlite")
    selected = preflight(source)
    job, _ = jobs.submit("owner", "routing-budget", source, selected, "192.168.22.8", launch=False)
    jobs.run(job["job_id"], "owner", source, selected, "192.168.22.8")
    result = jobs.get(job["job_id"], "owner")
    assert result["error"]["code"] == "job_time_budget_exhausted"
    assert result["error"]["details"]["failure_stage"] == "routing"


def test_watchdog_job_expiry_reaches_status_with_zero_remaining(tmp_path, monkeypatch):
    """The backend's shared-deadline error survives orchestration and persistence."""
    source = request(7)
    source["execution"]["max_seconds"] = 3600
    clock = [100.0]
    monkeypatch.setattr(suite_jobs.time, "monotonic", lambda: clock[0])
    monkeypatch.setattr(suite_backends, "switchyard_decision", lambda _: None)
    def expired(value, cancel, *, invocation, progress, job_deadline):
        assert value["execution"]["max_seconds"] == 1800
        clock[0] = job_deadline
        raise Problem("job_time_budget_exhausted", 504, failure_stage="subprocess")
    monkeypatch.setattr(suite_backends, "claude_generate", expired)
    jobs = suite_jobs.Jobs(tmp_path / "jobs.sqlite")
    selected = preflight(source)
    job, _ = jobs.submit("owner", "watchdog-budget", source, selected, "192.168.22.8", launch=False)
    jobs.run(job["job_id"], "owner", source, selected, "192.168.22.8")
    result = jobs.get(job["job_id"], "owner")
    assert result["status"] == "failed" and not jobs.results
    assert result["error"]["code"] == "job_time_budget_exhausted"
    assert result["error"]["details"]["failure_stage"] == "subprocess"
    assert result["execution_progress"]["job_remaining_seconds"] == 0