launcher/src/config/tests.rs

use crate::config::*;
use crate::pre_serve::*;
use std::collections::{HashMap, HashSet};

fn tcls() -> Cluster {
    Cluster {
        name: "test-cluster".to_string(),
        master_ib_ip: "192.168.1.1".to_string(),
        nodes: vec![
            Node {
                name: "node0".to_string(),
                ssh_ip: "10.0.0.1".to_string(),
                ib_ip: "192.168.1.1".to_string(),
                rank: 0,
                host_key: None,
                entity: None,
            },
            Node {
                name: "node1".to_string(),
                ssh_ip: "10.0.0.2".to_string(),
                ib_ip: "192.168.1.2".to_string(),
                rank: 1,
                host_key: None,
                entity: None,
            },
        ],
        extra_entities: vec![],
    }
}

fn tcfg() -> Config {
    Config {
        ssh_user: "u".to_string(),
        ssh_key_path: "/k".into(),
        clusters: vec![tcls()],
        nccl: HashMap::new(),
        hf_cache_dir: None,
        init_port: 25000,
        before_serve: HashMap::new(),
        exec_aliases: HashMap::new(),
        home_assistant: None,
        global_mods_dir: None,
        drop_caches: true,
        mount_aliases: HashMap::new(),
    }
}

fn tnccl() -> NcclConfig {
    NcclConfig {
        ib_hca: "rocep1s0f0".to_string(),
        socket_ifname: "enp1s0f0np0".to_string(),
        ib_gid_index: 3,
        ib_disable: 0,
        gloo_socket_ifname: None,
        ucx_net_devices: None,
        nccl_cross_nic: None,
        nccl_proto: None,
        nccl_p2p_level: None,
    }
}

#[test]
fn test_recipe_deserialize_minimal() {
    let r: Recipe = toml::from_str("\nname = \"M\"\nmodel = \"t/m\"\n").unwrap();
    assert_eq!(r.name, "M");
    assert_eq!(r.engine, "vllm");
    assert_eq!(r.port, 30000);
}

#[test]
fn test_recipe_deserialize_full() {
    let r: Recipe = toml::from_str("engine = \"sglang\"\nname = \"Full\"\ndocker_image = \"o/i:t\"\nmodel = \"o/m\"\nport = 30000\nrequired_nodes = 4\npipeline_parallel_size = 2\ndistributed_backend = \"ray\"\ncluster = \"c\"\nargs = [\"--a\", \"v\"]\n[env]\nK = \"V\"\n").unwrap();
    assert_eq!(r.engine, "sglang");
    assert_eq!(r.required_nodes, Some(4));
    assert_eq!(r.distributed_backend, Some(DistributedBackend::Ray));
}

#[test]
fn test_recipe_deserialize_no_required_nodes() {
    let r: Recipe = toml::from_str("name = \"F\"\nmodel = \"t/m\"\n").unwrap();
    assert_eq!(r.required_nodes, None);
    assert_eq!(r.pipeline_parallel_size, None);
}

#[test]
fn test_composed_source_exclude_deserialize() {
    let toml = r#"
name = "C"
model = "t/m"
[composed_model]
mount = "/root/model"
[[composed_model.sources]]
repo = "base"
[[composed_model.sources]]
repo = "overlay"
glob = "L*.safetensors"
exclude = ["L5.safetensors", "L8.safetensors"]
[[composed_model.files]]
name = "cfg.json"
content = "{}"
"#;
    let r: Recipe = toml::from_str(toml).unwrap();
    let composed = r.composed_model.unwrap();
    assert_eq!(composed.sources.len(), 2);
    let overlay = &composed.sources[1];
    assert_eq!(overlay.repo, "overlay");
    assert_eq!(overlay.glob.as_deref(), Some("L*.safetensors"));
    let exclude = overlay
        .exclude
        .as_ref()
        .expect("exclude should deserialize");
    assert_eq!(exclude.len(), 2);
    assert!(exclude.contains(&"L5.safetensors".to_string()));
    assert!(exclude.contains(&"L8.safetensors".to_string()));
}

#[test]
fn test_config_deserialize() {
    let c: Config = toml::from_str("ssh_user = \"a\"\nssh_key_path = \"/p\"\n[[clusters]]\nname = \"p\"\nmaster_ib_ip = \"1\"\nnodes = [{name=\"n1\",ssh_ip=\"10.0.0.10\",ib_ip=\"1\",rank=0},{name=\"n2\",ssh_ip=\"10.0.0.11\",ib_ip=\"2\",rank=1}]\n[nccl.\"p\"]\nib_hca = \"m\"\nsocket_ifname = \"e\"\nib_gid_index = 3\nib_disable = 0\n").unwrap();
    assert_eq!(c.clusters.len(), 1);
    assert_eq!(c.clusters[0].nodes.len(), 2);
}

#[test]
fn test_config_deserialize_with_host_key() {
    let c: Config = toml::from_str("ssh_user = \"a\"\nssh_key_path = \"/k\"\n[[clusters]]\nname = \"s\"\nmaster_ib_ip = \"1\"\nnodes = [\n    { name = \"n\", ssh_ip = \"2\", ib_ip = \"1\", rank = 0, host_key = \"SHA256:abc\" },\n]\n").unwrap();
    assert_eq!(c.clusters[0].nodes[0].host_key, Some("SHA256:abc".into()));
}

#[test]
fn test_select_cluster_by_name() {
    let mut cfg = tcfg();
    cfg.nccl.insert("test-cluster".into(), tnccl());
    let (c, n) = select_cluster(&cfg, "test-cluster", None, &HashSet::new()).unwrap();
    assert_eq!(c.name, "test-cluster");
    assert_eq!(n.ib_hca, "rocep1s0f0");
}

#[test]
fn test_select_cluster_auto() {
    let mut cfg = tcfg();
    cfg.nccl.insert("test-cluster".into(), tnccl());
    let (c, _n) = select_cluster(&cfg, "auto", None, &HashSet::new()).unwrap();
    assert_eq!(c.name, "test-cluster");
}

#[test]
fn test_select_cluster_auto_no_match() {
    assert!(select_cluster(&tcfg(), "auto", None, &HashSet::new()).is_err());
}

#[test]
fn test_select_cluster_not_found() {
    let mut cfg = tcfg();
    cfg.clusters = vec![];
    assert!(select_cluster(&cfg, "x", None, &HashSet::new()).is_err());
}

#[test]
fn test_select_cluster_by_name_found() {
    let mut nccl = HashMap::new();
    nccl.insert("test-cluster".into(), tnccl());
    let mut cfg = tcfg();
    cfg.nccl = nccl;
    let c = select_cluster_by_name(&cfg, "test-cluster").unwrap();
    assert_eq!(c.name, "test-cluster");
}

#[test]
fn test_select_cluster_by_name_not_found() {
    assert!(select_cluster_by_name(&tcfg(), "x").is_err());
}

#[test]
fn test_select_cluster_by_prefix() {
    let mut cfg = tcfg();
    cfg.nccl.insert("test-cluster".into(), tnccl());
    let c = select_cluster_by_name(&cfg, "test-c").unwrap();
    assert_eq!(c.name, "test-cluster");
    assert!(select_cluster_by_name(&cfg, "no-such").is_err());
}

#[test]
fn test_default_engine_is_vllm() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
    assert_eq!(r.engine, "vllm");
}

#[test]
fn test_default_distributed_backend_is_none() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
    assert_eq!(r.distributed_backend, None);
}

#[test]
fn test_deploy_defaults_to_docker() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
    assert_eq!(r.deploy, DeployTarget::Docker);
}

#[test]
fn test_deploy_explicit_tmux() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\ndeploy = \"tmux\"\n").unwrap();
    assert_eq!(r.deploy, DeployTarget::Tmux);
}

#[test]
fn test_deploy_explicit_docker() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\ndeploy = \"docker\"\n").unwrap();
    assert_eq!(r.deploy, DeployTarget::Docker);
}

#[test]
fn test_venv_field() {
    let r: Recipe = toml::from_str(
        "name = \"x\"\nmodel = \"m\"\ndeploy = \"tmux\"\nvenv = \"~/venvs/vllm025\"\n",
    )
    .unwrap();
    assert_eq!(r.deploy, DeployTarget::Tmux);
    assert_eq!(r.venv.as_deref(), Some("~/venvs/vllm025"));
}

#[test]
fn test_venv_defaults_to_none() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
    assert_eq!(r.venv, None);
}

#[test]
fn test_venv_setup_deserializes() {
    let r: Recipe = toml::from_str(
        r#"name = "x"
model = "m"
deploy = "tmux"
venv = "~/venvs/vllm025"
venv_setup = [
    "uv venv ~/venvs/vllm025 -p 3.12",
    "uv pip install vllm==0.25.1 --torch-backend=cu130",
]
"#,
    )
    .unwrap();
    assert_eq!(r.venv.as_deref(), Some("~/venvs/vllm025"));
    assert_eq!(r.venv_setup.len(), 2);
    assert_eq!(r.venv_setup[0], "uv venv ~/venvs/vllm025 -p 3.12");
    assert_eq!(
        r.venv_setup[1],
        "uv pip install vllm==0.25.1 --torch-backend=cu130"
    );
}

#[test]
fn test_venv_setup_defaults_to_empty() {
    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
    assert!(r.venv_setup.is_empty());
}

#[test]
fn default_config_dir_appends_sm121() {
    let Some(d) = default_config_dir() else {
        return;
    };
    assert!(d.ends_with("sm121") && d.parent().is_some());
}