//! 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) -> Config { Config { 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 { 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()); }