launcher/src/before_serve_tests.rs

//! Tests for `before_serve` (split at the 500-line cap).

use super::*;
use crate::config::{Node, PreServeCommand};
use crate::pre_serve::VenvSetup;
use std::collections::HashMap;

fn test_cluster() -> Cluster {
    Cluster {
        name: "c1".to_string(),
        master_ib_ip: Some("1.2.3.4".to_string()),
        nodes: vec![
            Node {
                name: "n0".to_string(),
                ssh_ip: "10.0.0.1".to_string(),
                ib_ip: Some("1.2.3.4".to_string()),
                rank: 0,
                host_key: None,
                entity: None,
            },
            Node {
                name: "n1".to_string(),
                ssh_ip: "10.0.0.2".to_string(),
                ib_ip: Some("1.2.3.5".to_string()),
                rank: 1,
                host_key: None,
                entity: None,
            },
        ],
        extra_entities: vec![],
    }
}

fn test_config(before_serve: HashMap<String, crate::config::BeforeServeTable>) -> Config {
    Config {
        cache: None,
        flare: None,
        ssh_user: "u".to_string(),
        ssh_key_path: "/k".into(),
        clusters: vec![test_cluster()],
        nccl: HashMap::new(),
        hf_cache_dir: None,
        init_port: 25000,
        before_serve,
        exec_aliases: HashMap::new(),
        home_assistant: None,
        mikrotik: None,
        global_mods_dir: 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,
    }
}

#[test]
fn truncate_keeps_short_string_unchanged() {
    assert_eq!(truncate("short", 120), "short");
    assert_eq!(truncate("", 120), "");
}

#[test]
fn truncate_shortens_to_char_boundary() {
    // 10 chars, cap at 4 → first 4 chars + ellipsis
    let out = truncate("abcdefghij", 4);
    assert_eq!(out, "abcd…");
}

#[test]
fn truncate_does_not_panic_on_multibyte() {
    // Each 'é' is 2 bytes; a byte slice `&s[..3]` would land mid-char and
    // panic. The char-based impl must not. (Regression for C-2.)
    let s = "éééééé"; // 6 chars, 12 bytes
    assert_eq!(truncate(s, 3), "ééé…");
    // cap exactly at a multibyte boundary must also be safe
    assert_eq!(truncate(s, 4), "éééé…");
}

#[test]
fn truncate_cjk_chars() {
    // 3-byte CJK: byte-slice at any non-multiple-of-3 would panic.
    let s = "你好世界你好";
    assert_eq!(truncate(s, 2), "你好…");
}

#[test]
fn combined_runs_global_then_cluster_then_recipe() {
    let mut bs = HashMap::new();
    bs.insert(
        "c1".to_string(),
        crate::config::BeforeServeTable {
            commands: vec![PreServeCommand::Bare("cluster_cmd".into())],
        },
    );
    bs.insert(
        "*".to_string(),
        crate::config::BeforeServeTable {
            commands: vec![PreServeCommand::Bare("global_cmd".into())],
        },
    );
    let cfg = test_config(bs);
    let recipe = Recipe {
        before_serve: vec![PreServeCommand::Bare("recipe_cmd".into())],
        ..test_recipe()
    };

    let cmds = combined_commands(&cfg, &test_cluster(), &recipe);
    let names: Vec<&str> = cmds.iter().map(|(_, c)| c.command()).collect();
    let origins: Vec<&str> = cmds.iter().map(|(o, _)| *o).collect();
    assert_eq!(names, vec!["global_cmd", "cluster_cmd", "recipe_cmd"]);
    assert_eq!(origins, vec!["global", "cluster", "recipe"]);
}

#[test]
fn combined_global_applies_to_unlisted_cluster() {
    let mut bs = HashMap::new();
    bs.insert(
        "*".to_string(),
        crate::config::BeforeServeTable {
            commands: vec![PreServeCommand::Bare("global_cmd".into())],
        },
    );
    // Cluster "c1" has no specific entry — only the global hook runs.
    let cfg = test_config(bs);
    let recipe = Recipe {
        before_serve: vec![],
        after_stop: vec![],
        solo: None,
        ..test_recipe()
    };

    let cmds = combined_commands(&cfg, &test_cluster(), &recipe);
    let names: Vec<&str> = cmds.iter().map(|(_, c)| c.command()).collect();
    assert_eq!(names, vec!["global_cmd"]);
}

#[test]
fn combined_orders_cluster_then_recipe() {
    let mut bs = HashMap::new();
    bs.insert(
        "c1".to_string(),
        crate::config::BeforeServeTable {
            commands: vec![PreServeCommand::Bare("cluster_cmd_1".into())],
        },
    );
    let cfg = test_config(bs);
    let recipe = Recipe {
        before_serve: vec![
            PreServeCommand::Bare("recipe_cmd_1".into()),
            PreServeCommand::Bare("recipe_cmd_2".into()),
        ],
        ..test_recipe()
    };

    let cmds = combined_commands(&cfg, &test_cluster(), &recipe);
    let names: Vec<&str> = cmds.iter().map(|(_, c)| c.command()).collect();
    let origins: Vec<&str> = cmds.iter().map(|(o, _)| *o).collect();
    assert_eq!(names, vec!["cluster_cmd_1", "recipe_cmd_1", "recipe_cmd_2"]);
    assert_eq!(origins, vec!["cluster", "recipe", "recipe"]);
}

#[test]
fn combined_empty_when_unconfigured() {
    let cfg = test_config(HashMap::new());
    let recipe = Recipe {
        before_serve: vec![],
        after_stop: vec![],
        solo: None,
        ..test_recipe()
    };
    assert!(combined_commands(&cfg, &test_cluster(), &recipe).is_empty());
}

fn test_recipe() -> Recipe {
    Recipe {
        flare: None,
        build_profile: None,
        parent_model: None,
        archived: false,
        engine: "vllm".to_string(),
        name: "Test".to_string(),
        docker_image: String::new(),
        model: "x".to_string(),
        port: 30000,
        required_nodes: None,
        pipeline_parallel_size: None,
        tensor_parallel_size: 1,
        args: crate::pre_serve::RecipeArgs::default(),
        cold_args: crate::pre_serve::RecipeArgs::default(),
        env: HashMap::new(),
        cluster: None,
        entrypoint: vec![],
        distributed_backend: Some(crate::config::DistributedBackend::Mp),
        before_serve: vec![],
        after_stop: vec![],
        solo: None,
        deploy: crate::config::DeployTarget::Docker,
        venv: None,
        venv_setup: VenvSetup::default(),
        venv_setup_script: None,
        mods: Default::default(),
        sync_image: true,
        docker_image_dir: None,
        shm_size: String::new(),
        devices: vec![],
        memory_limit: None,
        cpus: None,
        extra_mounts: vec![],
        extra_mounts_manifest: None,
        build: Default::default(),
        composed_model: None,
        staging: vec![],
        dgpp: None,
    }
}

// ── Deserialization tests (the ones that live in this module) ──

#[test]
fn parses_cluster_level_before_serve() {
    let toml_str = r#"
ssh_user = "u"
ssh_key_path = "/k"

[[clusters]]
name = "c1"
master_ib_ip = "1.1.1.1"
nodes = [
{ name = "n0", ssh_ip = "2.2.2.2", ib_ip = "1.1.1.1", rank = 0 },
]

[before_serve."c1"]
commands = [
"sudo nvidia-smi -lgc 300,2300",
{ command = "sudo nvidia-smi -pm 1", on_failure = "warn" },
]
"#;
    let cfg: Config = toml::from_str(toml_str).unwrap();
    assert!(cfg.before_serve.contains_key("c1"));
    let list = &cfg.before_serve["c1"].commands;
    assert_eq!(list.len(), 2);
    assert_eq!(list[0].command(), "sudo nvidia-smi -lgc 300,2300");
    assert_eq!(list[0].on_failure(), OnFailure::Abort); // default
    assert_eq!(list[1].command(), "sudo nvidia-smi -pm 1");
    assert_eq!(list[1].on_failure(), OnFailure::Warn);
}

#[test]
fn parses_recipe_level_before_serve() {
    let toml_str = r#"
name = "Demo"
model = "x"
before_serve = [
"echo a",
{ command = "echo b", on_failure = "continue" },
]
"#;
    let r: Recipe = toml::from_str(toml_str).unwrap();
    assert_eq!(r.before_serve.len(), 2);
    assert_eq!(r.before_serve[0].command(), "echo a");
    assert_eq!(r.before_serve[0].on_failure(), OnFailure::Abort);
    assert_eq!(r.before_serve[1].command(), "echo b");
    assert_eq!(r.before_serve[1].on_failure(), OnFailure::Continue);
}

/// A recipe-relative hook script is read from the recipe dir; absolute
/// paths and `..` traversal are refused (a synced recipe repo must not be
/// able to run an arbitrary file on a node).
#[tokio::test]
async fn recipe_scripts_resolve_inside_the_recipe_dir() {
    let dir = std::env::temp_dir().join(format!("sm121-hook-{}", std::process::id()));
    tokio::fs::create_dir_all(dir.join("scripts"))
        .await
        .unwrap();
    tokio::fs::write(dir.join("scripts/prepare.sh"), "echo hi\n")
        .await
        .unwrap();

    let body = read_recipe_script("scripts/prepare.sh", Some(&dir), "n0")
        .await
        .expect("recipe script reads");
    assert_eq!(body, "echo hi\n");

    let abs = read_recipe_script("/etc/passwd", Some(&dir), "n0").await;
    assert!(abs.is_err(), "absolute paths are refused");
    let trav = read_recipe_script("scripts/../../etc/passwd", Some(&dir), "n0").await;
    assert!(trav.is_err(), "`..` traversal is refused");
    let missing = read_recipe_script("scripts/nope.sh", Some(&dir), "n0").await;
    assert!(missing.is_err(), "missing file names the built-in aliases");
    let no_dir = read_recipe_script("scripts/prepare.sh", None, "n0").await;
    assert!(no_dir.is_err(), "unresolvable recipe dir is an error");

    let _ = tokio::fs::remove_dir_all(&dir).await;
}

/// A built-in alias wins over a same-named file in recipe sources.
#[test]
fn builtin_aliases_resolve_before_recipe_files() {
    assert!(crate::builtin_scripts::resolve("stop_dgx_services").is_some());
    assert!(crate::builtin_scripts::resolve("scripts/prepare.sh").is_none());
}