from scripts.quant.recipes import METHOD_SPECS, method_names, vllm_serve_args


def test_all_seven_methods_present():
    assert method_names() == ["bf16", "fp8", "int8", "awq", "gptq", "bnb8", "bnb4"]


def test_t4_flags():
    assert METHOD_SPECS["fp8"]["t4_deployable"] is False
    assert METHOD_SPECS["awq"]["t4_deployable"] is True
    assert METHOD_SPECS["bnb8"]["t4_deployable"] is True


def test_baseline_serves_base_model():
    args = vllm_serve_args("bf16", model_path="/m/base", base_model="/m/base", port=8001)
    assert "--model" in args and "/m/base" in args
    assert "--quantization" not in args
    assert args[args.index("--port") + 1] == "8001"


def test_fp8_uses_online_quantization_on_base():
    args = vllm_serve_args("fp8", model_path="/m/base", base_model="/m/base", port=8001)
    assert args[args.index("--quantization") + 1] == "fp8"
    assert "/m/base" in args


def test_compressor_serves_built_path_without_quant_flag():
    args = vllm_serve_args("awq", model_path="/m/awq", base_model="/m/base", port=8001)
    assert "/m/awq" in args
    assert "--quantization" not in args  # quant config travels with the compressed checkpoint


def test_bnb_uses_bitsandbytes_flags():
    args = vllm_serve_args("bnb4", model_path="/m/bnb4", base_model="/m/base", port=8001)
    assert args[args.index("--quantization") + 1] == "bitsandbytes"
    assert args[args.index("--load-format") + 1] == "bitsandbytes"
    assert "/m/bnb4" in args