launcher/src/power/tests.rs

use super::*;
use crate::power::toggle::off_extras;

fn cfg_with_ha() -> Config {
    let toml = r#"
ssh_user = "u"
ssh_key_path = "/k"
[[clusters]]
name = "x"
master_ib_ip = "1.1.1.1"
nodes = [
    { name = "n1", ssh_ip = "10.0.0.1", ib_ip = "10.0.0.1", rank = 0, entity = "switch.n1" },
    { name = "n2", ssh_ip = "10.0.0.2", ib_ip = "10.0.0.2", rank = 1 },
]
extra_entities = ["switch.network_x"]
[nccl."x"]
ib_hca = "h"
socket_ifname = "s"
ib_gid_index = 1
ib_disable = 0

[home_assistant]
url = "http://ha.local:8123"
token = "abcdef0123456789abcdef0123456789"
shutdown_delay_seconds = 5
"#;
    toml::from_str(toml).unwrap()
}

fn cfg_no_ha() -> Config {
    let toml = r#"
ssh_user = "u"
ssh_key_path = "/k"
[[clusters]]
name = "x"
master_ib_ip = "1.1.1.1"
nodes = [
    { name = "n1", ssh_ip = "10.0.0.1", ib_ip = "10.0.0.1", rank = 0, entity = "switch.n1" },
]
[nccl."x"]
ib_hca = "h"
socket_ifname = "s"
ib_gid_index = 1
ib_disable = 0
"#;
    toml::from_str(toml).unwrap()
}

#[test]
fn select_targets_filters_by_node() {
    let cfg = cfg_with_ha();
    let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
    let t = select_targets(&c, &["n1".to_string()], "test").unwrap();
    assert_eq!(t.len(), 1);
    assert_eq!(t[0].name, "n1");
}

#[test]
fn select_targets_all() {
    let cfg = cfg_with_ha();
    let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
    let t = select_targets(&c, &[], "test").unwrap();
    assert_eq!(t.len(), 2);
}

#[test]
fn select_targets_multiple_nodes() {
    // `stop --node a --node b --off` must power off *both* nodes, not just
    // the first. (Regression for I-5.)
    let cfg = cfg_with_ha();
    let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
    let t = select_targets(&c, &["n1".to_string(), "n2".to_string()], "test").unwrap();
    assert_eq!(t.len(), 2);
    let names: Vec<_> = t.iter().map(|n| n.name.as_str()).collect();
    assert!(names.contains(&"n1"));
    assert!(names.contains(&"n2"));
}

#[test]
fn select_targets_partial_missing_errors() {
    let cfg = cfg_with_ha();
    let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
    // One valid + one missing → error listing the missing one.
    let r = select_targets(&c, &["n1".to_string(), "nope".to_string()], "test");
    assert!(r.is_err());
    assert!(r.unwrap_err().to_string().contains("nope"));
}

#[test]
fn select_targets_unknown_node_errors() {
    let cfg = cfg_with_ha();
    let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
    let r = select_targets(&c, &["nope".to_string()], "test");
    assert!(r.is_err());
}

#[test]
fn power_on_requires_ha_config() {
    let cfg = cfg_no_ha();
    assert!(cfg.home_assistant.is_none());
}

#[test]
fn cluster_extra_entities_parsed() {
    use crate::config::ExtraEntity;
    let cfg = cfg_with_ha();
    let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
    assert_eq!(
        c.extra_entities,
        vec![ExtraEntity::Simple("switch.network_x".into())]
    );
}

#[test]
fn off_extras_whole_cluster_keeps_extras_subset_skips_them() {
    use std::collections::HashSet;
    let empty: HashSet<String> = HashSet::new();
    // A whole-cluster power off (no --node) turns the shared fabric off too.
    let cfg = cfg_with_ha();
    let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
    assert_eq!(
        off_extras(&c, &[], &empty),
        vec!["switch.network_x".to_string()]
    );
    // A subset power off (--node) must NOT take the shared switch down —
    // the nodes that stay up depend on it (NCCL over the fabric).
    assert_eq!(
        off_extras(&c, &["n1".to_string()], &empty),
        Vec::<String>::new()
    );
    assert_eq!(
        off_extras(&c, &["n1".to_string(), "n2".to_string()], &empty),
        Vec::<String>::new()
    );
}

#[test]
fn off_extras_whole_cluster_keeps_extra_held_by_other_active_cluster() {
    // The fabric switch is shared across clusters. A whole-cluster off of `x`
    // must NOT turn off `switch.network_x` while another cluster that also
    // lists it has an open run — that serve depends on the same switch.
    use std::collections::HashSet;
    let cfg = cfg_with_ha();
    let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
    let held: HashSet<String> = ["switch.network_x".to_string()].into_iter().collect();
    // held → the shared switch is kept on, nothing is turned off
    assert_eq!(off_extras(&c, &[], &held), Vec::<String>::new());
    // an unrelated held entity does not protect this cluster's switch
    let other: HashSet<String> = ["switch.somewhere_else".to_string()].into_iter().collect();
    assert_eq!(
        off_extras(&c, &[], &other),
        vec!["switch.network_x".to_string()]
    );
}

#[test]
fn extra_entity_forms_blocking_and_soft() {
    use crate::config::ExtraEntity;
    let toml = r#"
ssh_user = "u"
ssh_key_path = "/k"
[[clusters]]
name = "x"
nodes = [{ name = "n1", ssh_ip = "10.0.0.1", rank = 0 }]
extra_entities = ["switch.plug", { entity = "switch.fan", blocks = false }, { entity = "switch.other" }]
"#;
    let cfg: crate::config::Config = toml::from_str(toml).unwrap();
    let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
    let e = &c.extra_entities;
    assert_eq!(e.len(), 3);
    assert!(matches!(&e[0], ExtraEntity::Simple(s) if s == "switch.plug"));
    assert!(e[0].blocks());
    assert!(
        matches!(&e[1], ExtraEntity::Detailed { entity, blocks } if entity == "switch.fan" && !blocks)
    );
    assert!(!e[1].blocks());
    // table form without `blocks` defaults to blocking
    assert!(
        matches!(&e[2], ExtraEntity::Detailed { entity, blocks } if entity == "switch.other" && *blocks)
    );
    assert!(e[2].blocks());
    assert_eq!(e[1].entity(), "switch.fan");
}

#[test]
fn cluster_without_extra_entities_is_empty() {
    let cfg = cfg_no_ha();
    let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
    assert!(c.extra_entities.is_empty());
}

#[test]
fn mikrotik_dispatch_lookup() {
    let cfg = cfg_with_ha();
    // No [mikrotik] in this config — every entity stays on the HA path.
    assert!(super::mikro::device(cfg.mikrotik.as_ref(), "switch.n1").is_none());

    let toml = r#"
ssh_user = "u"
ssh_key_path = "/k"
[[clusters]]
name = "x"
nodes = [{ name = "n1", ssh_ip = "10.0.0.1", rank = 0, entity = "router.lab" }]
extra_entities = ["router.lab", "switch.plug"]
[mikrotik]
[[mikrotik.devices]]
name = "router.lab"
address = "192.168.88.1"
username = "admin"
password = "p"
"#;
    let cfg: crate::config::Config = toml::from_str(toml).unwrap();
    let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
    let off = super::mikro::collect_off(&cfg, &c.nodes, &super::extra_ids(&c));
    assert_eq!(off, vec!["router.lab".to_string()]);
    // same device referenced twice (node + extra) → listed once
    let dev = super::mikro::device(cfg.mikrotik.as_ref(), "router.lab");
    assert!(dev.is_some());
}

#[test]
fn decide_on_cycles_only_when_plug_on_and_node_down() {
    use crate::power::OnAction;
    use crate::power::decide_on;
    assert_eq!(decide_on(Some("on"), true), OnAction::Cycle);
    assert_eq!(decide_on(Some("on"), false), OnAction::AlreadyOn);
    assert_eq!(decide_on(Some("off"), true), OnAction::TurnOn);
    assert_eq!(decide_on(None, true), OnAction::TurnOn);
}