launcher/src/cluster_state/tests.rs

//! Tests for [`crate::cluster_state`] (separate file — the main
//! module sits at the 500-line cap).

use crate::cluster_state::*;
#[cfg(test)]
use crate::config::{Config, NcclConfig, Node};
use std::collections::HashMap;

fn test_node(name: &str) -> Node {
    Node {
        name: name.to_string(),
        ssh_ip: format!("10.0.0.{}", name.chars().last().unwrap_or('0')),
        ib_ip: Some("10.0.1.1".to_string()),
        rank: 0,
        host_key: None,
        entity: None,
    }
}

fn test_cluster(name: &str, node_names: &[&str]) -> Cluster {
    Cluster {
        name: name.to_string(),
        master_ib_ip: Some("10.0.1.1".to_string()),
        nodes: node_names.iter().map(|n| test_node(n)).collect(),
        extra_entities: vec![],
    }
}

fn test_nccl() -> NcclConfig {
    NcclConfig {
        ib_hca: "r".to_string(),
        socket_ifname: "e".to_string(),
        ib_gid_index: Some(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 verdict_all_up() {
    let v = verdict(&[(true, false), (true, true)], false);
    assert!(v.up);
    assert!(!v.switch_off);
    assert_eq!(v.free, 1);
}

#[test]
fn verdict_one_offline_node_takes_cluster_down() {
    let v = verdict(&[(true, false), (false, false)], false);
    assert!(!v.up);
    assert_eq!(v.free, 0);
}

#[test]
fn verdict_switch_off_takes_cluster_down() {
    let v = verdict(&[(true, false), (true, false)], true);
    assert!(!v.up);
    assert!(v.switch_off);
    assert_eq!(v.free, 0);
}

#[test]
fn free_count_all_reachable() {
    let mut state = ClusterState::default();
    state.nodes.insert(
        "n1".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );
    state.nodes.insert(
        "n2".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );
    let c = test_cluster("c", &["n1", "n2"]);
    assert_eq!(state.free_count(&c), 2);
}

#[test]
fn free_count_occupied_excluded() {
    let mut state = ClusterState::default();
    state.nodes.insert(
        "n1".into(),
        NodeState {
            reachable: true,
            occupied: true,
        },
    );
    state.nodes.insert(
        "n2".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );
    let c = test_cluster("c", &["n1", "n2"]);
    assert_eq!(state.free_count(&c), 1);
}

#[test]
fn free_count_unreachable_excluded() {
    let mut state = ClusterState::default();
    state.nodes.insert(
        "n1".into(),
        NodeState {
            reachable: false,
            occupied: false,
        },
    );
    state.nodes.insert(
        "n2".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );
    let c = test_cluster("c", &["n1", "n2"]);
    assert_eq!(state.free_count(&c), 1);
}

#[test]
fn free_count_missing_node_not_free() {
    let mut state = ClusterState::default();
    state.nodes.insert(
        "n1".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );
    // n2 not probed — absent is treated as not free
    let c = test_cluster("c", &["n1", "n2"]);
    assert_eq!(state.free_count(&c), 1);
}

#[test]
fn select_smallest_with_enough_free() {
    let mut state = ClusterState::default();
    // big cluster: 4 nodes, 2 free
    state.nodes.insert(
        "n1".into(),
        NodeState {
            reachable: true,
            occupied: true,
        },
    );
    state.nodes.insert(
        "n2".into(),
        NodeState {
            reachable: true,
            occupied: true,
        },
    );
    state.nodes.insert(
        "n3".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );
    state.nodes.insert(
        "n4".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );
    // small cluster: 2 nodes, 1 free
    state.nodes.insert(
        "n5".into(),
        NodeState {
            reachable: true,
            occupied: true,
        },
    );
    state.nodes.insert(
        "n6".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );

    let cfg = Config {
        ssh_user: "u".to_string(),
        ssh_key_path: "/k".into(),
        clusters: vec![
            test_cluster("big", &["n1", "n2", "n3", "n4"]),
            test_cluster("small", &["n5", "n6"]),
        ],
        nccl: {
            let mut m = HashMap::new();
            m.insert("big".into(), test_nccl());
            m.insert("small".into(), test_nccl());
            m
        },
        hf_cache_dir: None,
        init_port: 25000,
        before_serve: HashMap::new(),
        exec_aliases: HashMap::new(),
        home_assistant: None,
        mikrotik: None,
        drop_caches: true,
        mount_aliases: HashMap::new(),
        recipe_sources: Vec::new(),
        recipe_repos: Vec::new(),
        docker_image_sources: Vec::new(),
        cluster_sources: Vec::new(),
        registry: None,
        global_mods_dir: None,
    };

    // Request 1 node — small cluster is smallest with free >= 1
    let (c, _) = select_cluster(&cfg, "auto", Some(1), &state).unwrap();
    assert_eq!(c.name, "small");

    // Request 2 nodes — only big cluster has 2 free
    let (c, _) = select_cluster(&cfg, "auto", Some(2), &state).unwrap();
    assert_eq!(c.name, "big");
}

#[test]
fn select_cluster_not_found_errors() {
    let state = ClusterState::default();
    let cfg = Config {
        ssh_user: "u".to_string(),
        ssh_key_path: "/k".into(),
        clusters: vec![],
        nccl: HashMap::new(),
        hf_cache_dir: None,
        init_port: 25000,
        before_serve: HashMap::new(),
        exec_aliases: HashMap::new(),
        home_assistant: None,
        mikrotik: None,
        drop_caches: true,
        mount_aliases: HashMap::new(),
        recipe_sources: Vec::new(),
        recipe_repos: Vec::new(),
        docker_image_sources: Vec::new(),
        cluster_sources: Vec::new(),
        registry: None,
        global_mods_dir: None,
    };
    assert!(select_cluster(&cfg, "nonexistent", None, &state).is_err());
}

#[test]
fn no_cluster_has_enough_free_errors() {
    let mut state = ClusterState::default();
    state.nodes.insert(
        "n1".into(),
        NodeState {
            reachable: true,
            occupied: true,
        },
    );
    let cfg = Config {
        ssh_user: "u".to_string(),
        ssh_key_path: "/k".into(),
        clusters: vec![test_cluster("c", &["n1"])],
        nccl: {
            let mut m = HashMap::new();
            m.insert("c".into(), test_nccl());
            m
        },
        hf_cache_dir: None,
        init_port: 25000,
        before_serve: HashMap::new(),
        exec_aliases: HashMap::new(),
        home_assistant: None,
        mikrotik: None,
        drop_caches: true,
        mount_aliases: HashMap::new(),
        recipe_sources: Vec::new(),
        recipe_repos: Vec::new(),
        docker_image_sources: Vec::new(),
        cluster_sources: Vec::new(),
        registry: None,
        global_mods_dir: None,
    };
    assert!(select_cluster(&cfg, "auto", Some(1), &state).is_err());
}

#[test]
fn single_node_allows_cluster_without_nccl() {
    let mut state = ClusterState::default();
    state.nodes.insert(
        "n1".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );
    let cfg = Config {
        ssh_user: "u".to_string(),
        ssh_key_path: "/k".into(),
        clusters: vec![test_cluster("c", &["n1"])],
        nccl: HashMap::new(), // no NCCL config
        hf_cache_dir: None,
        init_port: 25000,
        before_serve: HashMap::new(),
        exec_aliases: HashMap::new(),
        home_assistant: None,
        mikrotik: None,
        drop_caches: true,
        mount_aliases: HashMap::new(),
        recipe_sources: Vec::new(),
        recipe_repos: Vec::new(),
        docker_image_sources: Vec::new(),
        cluster_sources: Vec::new(),
        registry: None,
        global_mods_dir: None,
    };
    // Single-node recipe on a cluster without NCCL config → OK
    let (c, nccl) = select_cluster(&cfg, "auto", Some(1), &state).unwrap();
    assert_eq!(c.name, "c");
    assert!(nccl.is_none());
}

#[test]
fn multi_node_requires_nccl() {
    let mut state = ClusterState::default();
    state.nodes.insert(
        "n1".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );
    state.nodes.insert(
        "n2".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );
    let cfg = Config {
        ssh_user: "u".to_string(),
        ssh_key_path: "/k".into(),
        clusters: vec![test_cluster("c", &["n1", "n2"])],
        nccl: HashMap::new(), // no NCCL config
        hf_cache_dir: None,
        init_port: 25000,
        before_serve: HashMap::new(),
        exec_aliases: HashMap::new(),
        home_assistant: None,
        mikrotik: None,
        drop_caches: true,
        mount_aliases: HashMap::new(),
        recipe_sources: Vec::new(),
        recipe_repos: Vec::new(),
        docker_image_sources: Vec::new(),
        cluster_sources: Vec::new(),
        registry: None,
        global_mods_dir: None,
    };
    // Multi-node recipe on a cluster without NCCL config → error
    assert!(select_cluster(&cfg, "auto", Some(2), &state).is_err());
}