from inferscale.models import SimulationConfig
from inferscale.optimizer import capacity_search, compare_schedulers, evaluate_rate


def test_capacity_search_returns_trace():
    cfg = {
        "duration_s": 6,
        "prompt_tokens_mean": 128,
        "output_tokens_mean": 16,
        "slo_ttft_ms": 1000,
        "slo_e2e_ms": 10000,
        "slo_attainment_target": 0.95,
    }
    out = capacity_search(cfg, min_rate=0.25, max_rate=3, iterations=3, repetitions=1)
    assert out["trace"]
    assert out["capacity_rps"] >= 0
    assert out["criterion"] == "all-repetitions-meet-target-and-drain"


def test_rate_evaluation_exposes_repetition_range():
    out = evaluate_rate(SimulationConfig(duration_s=4, request_rate_rps=1), 1.0, repetitions=2)
    assert out["slo_attainment_min"] <= out["slo_attainment"] <= out["slo_attainment_max"]
    assert out["repetitions"] == 2


def test_scheduler_compare_has_rows_and_diagnostics():
    out = compare_schedulers({"duration_s": 4, "request_rate_rps": 1, "output_tokens_mean": 8})
    assert len(out["rows"]) == 5
    assert all(row["bottleneck"] for row in out["rows"])