launcher/src/orchestrate/tests.rs

use super::*;
#[path = "retry_tests.rs"]
mod retry_tests;
use crate::cluster_state::NodeState;
use crate::config::{NcclConfig, Node};
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};

// --- Mock ends ---

struct MockProber {
    state: ClusterState,
}
impl NodeProber for MockProber {
    async fn probe(&self, _config: &Config) -> ClusterState {
        self.state.clone()
    }
}

struct MockPower {
    results: HashMap<String, bool>,
    turned_on: std::sync::Mutex<Vec<String>>,
}
impl PowerController for MockPower {
    async fn entity_is_on(&self, entity: &str) -> Option<bool> {
        self.results.get(entity).copied()
    }
    async fn turn_on(&self, entity: &str) -> Result<()> {
        self.turned_on.lock().unwrap().push(entity.to_string());
        Ok(())
    }
}

struct MockPinger {
    reachable: bool,
}
impl NodePinger for MockPinger {
    async fn is_reachable(&self, _config: &Config, _node: &Node) -> bool {
        self.reachable
    }
}

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: Some(format!("switch.{}", name)),
    }
}

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,
    }
}

fn test_config(clusters: Vec<Cluster>) -> Config {
    let mut nccl = HashMap::new();
    for c in &clusters {
        nccl.insert(c.name.clone(), test_nccl());
    }
    Config {
        ssh_user: "u".to_string(),
        ssh_key_path: "/k".into(),
        clusters,
        nccl,
        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(),
        docker_image_sources: Vec::new(),
        cluster_sources: Vec::new(),
        registry: None,
        global_mods_dir: None,
    }
}

fn test_recipe(required_nodes: Option<u32>) -> Recipe {
    toml::from_str(&format!(
        r#"
name = "test"
model = "test/model"
required_nodes = {}
"#,
        required_nodes.map_or("".into(), |n| n.to_string())
    ))
    .unwrap()
}

#[tokio::test]
async fn cluster_off_auto_power_on_powers_and_selects() {
    let c = test_cluster("c", &["n1", "n2"]);
    let config = test_config(vec![c.clone()]);
    let recipe = test_recipe(Some(2));

    let mut state = ClusterState::default();
    state.nodes.insert(
        "n1".into(),
        NodeState {
            reachable: false,
            occupied: false,
        },
    );
    state.nodes.insert(
        "n2".into(),
        NodeState {
            reachable: false,
            occupied: false,
        },
    );

    let mut after = ClusterState::default();
    after.nodes.insert(
        "n1".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );
    after.nodes.insert(
        "n2".into(),
        NodeState {
            reachable: true,
            occupied: false,
        },
    );

    let power = MockPower {
        results: [("switch.n1".into(), false), ("switch.n2".into(), false)]
            .into_iter()
            .collect(),
        turned_on: std::sync::Mutex::new(Vec::new()),
    };

    let call_count = AtomicUsize::new(0);
    struct CountProber {
        first: ClusterState,
        second: ClusterState,
        count: AtomicUsize,
    }
    impl NodeProber for CountProber {
        async fn probe(&self, _config: &Config) -> ClusterState {
            let n = self.count.fetch_add(1, Ordering::SeqCst);
            if n == 0 {
                self.first.clone()
            } else {
                self.second.clone()
            }
        }
    }
    let prober = CountProber {
        first: state,
        second: after,
        count: call_count,
    };

    let pinger = MockPinger { reachable: true };

    let (cluster, _nccl) = prepare_cluster(
        &config,
        &recipe,
        "auto",
        None,
        &PrepareRequest {
            prober: &prober,
            power: Some(&power),
            pinger: &pinger,
            auto_power_on: true,
            boot_timing: BootTiming {
                boot_wait_secs: 0,
                poll_interval_secs: 0,
                max_poll_secs: 10,
                retry_grace_secs: 0,
            },
            force_occupied: false,
        },
    )
    .await
    .unwrap();

    assert_eq!(cluster.name, "c");
    assert_eq!(cluster.nodes.len(), 2);
    let turned_on = power.turned_on.lock().unwrap();
    assert!(turned_on.contains(&"switch.n1".to_string()));
    assert!(turned_on.contains(&"switch.n2".to_string()));
}

#[tokio::test]
async fn cluster_off_no_auto_power_fails() {
    let c = test_cluster("c", &["n1", "n2"]);
    let config = test_config(vec![c.clone()]);
    let recipe = test_recipe(Some(2));

    let mut state = ClusterState::default();
    state.nodes.insert(
        "n1".into(),
        NodeState {
            reachable: false,
            occupied: false,
        },
    );
    state.nodes.insert(
        "n2".into(),
        NodeState {
            reachable: false,
            occupied: false,
        },
    );

    let prober = MockProber { state };
    let pinger = MockPinger { reachable: false };
    let power = MockPower {
        results: HashMap::new(),
        turned_on: std::sync::Mutex::new(Vec::new()),
    };

    let result = prepare_cluster(
        &config,
        &recipe,
        "auto",
        None,
        &PrepareRequest {
            prober: &prober,
            power: Some(&power),
            pinger: &pinger,
            auto_power_on: false,
            boot_timing: BootTiming {
                boot_wait_secs: 0,
                poll_interval_secs: 0,
                max_poll_secs: 10,
                retry_grace_secs: 0,
            },
            force_occupied: false,
        },
    )
    .await;

    assert!(result.is_err());
    let err = result.unwrap_err().to_string();
    assert!(err.contains("No suitable cluster found") || err.contains("requires"));
}

#[tokio::test]
async fn cluster_already_on_selects_directly() {
    let c = test_cluster("c", &["n1", "n2"]);
    let config = test_config(vec![c.clone()]);
    let recipe = test_recipe(Some(2));

    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 prober = MockProber { state };
    let pinger = MockPinger { reachable: true };
    let power = MockPower {
        results: [("switch.n1".into(), true), ("switch.n2".into(), true)]
            .into_iter()
            .collect(),
        turned_on: std::sync::Mutex::new(Vec::new()),
    };

    let (cluster, _nccl) = prepare_cluster(
        &config,
        &recipe,
        "auto",
        None,
        &PrepareRequest {
            prober: &prober,
            power: Some(&power),
            pinger: &pinger,
            auto_power_on: true,
            boot_timing: BootTiming {
                boot_wait_secs: 0,
                poll_interval_secs: 0,
                max_poll_secs: 10,
                retry_grace_secs: 0,
            },
            force_occupied: false,
        },
    )
    .await
    .unwrap();

    assert_eq!(cluster.name, "c");
    assert_eq!(cluster.nodes.len(), 2);
    assert!(power.turned_on.lock().unwrap().is_empty());
}

#[tokio::test]
async fn occupied_nodes_excluded() {
    let c = test_cluster("c", &["n1", "n2"]);
    let config = test_config(vec![c.clone()]);
    let recipe = test_recipe(Some(1));

    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 prober = MockProber { state };
    let pinger = MockPinger { reachable: true };
    let power = MockPower {
        results: HashMap::new(),
        turned_on: std::sync::Mutex::new(Vec::new()),
    };

    let (cluster, _nccl) = prepare_cluster(
        &config,
        &recipe,
        "auto",
        None,
        &PrepareRequest {
            prober: &prober,
            power: Some(&power),
            pinger: &pinger,
            auto_power_on: false,
            boot_timing: BootTiming {
                boot_wait_secs: 0,
                poll_interval_secs: 0,
                max_poll_secs: 10,
                retry_grace_secs: 0,
            },
            force_occupied: false,
        },
    )
    .await
    .unwrap();

    assert_eq!(cluster.nodes.len(), 1);
    assert_eq!(cluster.nodes[0].name, "n2");
}

#[tokio::test]
async fn force_occupied_keeps_occupied_nodes() {
    // `--force` skips the occupancy filter: an occupied node is kept in the
    // available set (reachable is still required). Mirror of
    // `occupied_nodes_excluded` with force_occupied: true.
    let c = test_cluster("c", &["n1", "n2"]);
    let config = test_config(vec![c.clone()]);
    let recipe = test_recipe(Some(2));

    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 prober = MockProber { state };
    let pinger = MockPinger { reachable: true };
    let power = MockPower {
        results: HashMap::new(),
        turned_on: std::sync::Mutex::new(Vec::new()),
    };

    let (cluster, _nccl) = prepare_cluster(
        &config,
        &recipe,
        "auto",
        None,
        &PrepareRequest {
            prober: &prober,
            power: Some(&power),
            pinger: &pinger,
            auto_power_on: false,
            boot_timing: BootTiming {
                boot_wait_secs: 0,
                poll_interval_secs: 0,
                max_poll_secs: 10,
                retry_grace_secs: 0,
            },
            force_occupied: true,
        },
    )
    .await
    .unwrap();

    // Both nodes kept (n1 occupied but force_occupied=true).
    assert_eq!(cluster.nodes.len(), 2);
}

#[test]
fn dry_run_resolves_without_probing_when_all_nodes_unreachable() {
    // Dry-run must resolve a cluster even when no node is reachable — the
    // whole point is to preview the command without touching hardware.
    let c = test_cluster("c", &["n1", "n2"]);
    let config = test_config(vec![c.clone()]);
    let recipe = test_recipe(Some(2));

    let (cluster, _nccl) =
        crate::orchestrate::prepare_cluster_dry_run(&config, &recipe, "auto", None).unwrap();
    assert_eq!(cluster.name, "c");
    // Dry-run keeps every configured node (no free-node filtering).
    assert_eq!(cluster.nodes.len(), 2);
}

#[test]
fn dry_run_solo_node_skips_probing() {
    let c = test_cluster("c", &["n1", "n2"]);
    let config = test_config(vec![c.clone()]);
    let recipe = test_recipe(Some(1));

    let (cluster, _nccl) =
        crate::orchestrate::prepare_cluster_dry_run(&config, &recipe, "auto", Some("n2")).unwrap();
    assert_eq!(cluster.nodes.len(), 1);
    assert_eq!(cluster.nodes[0].name, "n2");
}