from inferscale import compare_topologies, design_space_search


def test_topology_comparison_has_four_scenarios():
    result = compare_topologies({
        "model": "Qwen2.5-3B",
        "accelerator": "L4",
        "prefill_accelerator": "L4",
        "decode_accelerator": "L4",
        "quantization": "int8",
        "duration_s": 3,
        "request_rate_rps": 1,
        "prompt_tokens_mean": 64,
        "output_tokens_mean": 4,
        "prefix_reuse_fraction": 0.5,
        "shared_prefix_tokens": 32,
    })
    assert len(result["rows"]) == 4
    assert {row["topology"] for row in result["rows"]} == {"colocated", "disaggregated_pd"}
    assert all(row["goodput_per_accelerator"] >= 0 for row in result["rows"])


def test_design_space_marks_pareto_candidates():
    result = design_space_search({
        "model": "Qwen2.5-3B",
        "accelerator": "L4",
        "prefill_accelerator": "L4",
        "decode_accelerator": "L4",
        "quantization": "int8",
        "duration_s": 2,
        "request_rate_rps": 0.8,
        "prompt_tokens_mean": 64,
        "output_tokens_mean": 4,
        "shared_prefix_tokens": 32,
        "prefix_reuse_fraction": 0.5,
    }, include_disaggregated=False)
    assert result["candidate_count"] == 10
    assert result["pareto_count"] >= 1
    assert result["efficiency_pareto_count"] >= 1
    assert any(row["pareto"] for row in result["rows"])
    assert any(row["efficiency_pareto"] for row in result["rows"])