consortium_ansible/
inventory.rs

1//! Generate ansible inventory from FleetConfig.
2
3use std::collections::HashMap;
4
5use consortium_nix::config::FleetConfig;
6
7/// Generate an ansible inventory YAML string from fleet config.
8pub fn generate_inventory(config: &FleetConfig) -> String {
9    let ansible = config.ansible_config.as_ref();
10
11    let mut groups: HashMap<String, Vec<String>> = HashMap::new();
12
13    // Group by tags
14    for (name, node) in &config.nodes {
15        for tag in &node.tags {
16            groups.entry(tag.clone()).or_default().push(name.clone());
17        }
18        // All hosts go in the "all" group
19        groups
20            .entry("all".to_string())
21            .or_default()
22            .push(name.clone());
23    }
24
25    // Add custom host groups from ansible config
26    if let Some(ac) = ansible {
27        for (group, hosts) in &ac.host_groups {
28            groups
29                .entry(group.clone())
30                .or_default()
31                .extend(hosts.clone());
32        }
33    }
34
35    // Build YAML
36    let mut yaml = String::from("all:\n  children:\n");
37    for (group, hosts) in &groups {
38        if group == "all" {
39            continue;
40        }
41        yaml.push_str(&format!("    {}:\n      hosts:\n", group));
42        for host in hosts {
43            if let Some(node) = config.nodes.get(host) {
44                yaml.push_str(&format!(
45                    "        {}:\n          ansible_host: {}\n          ansible_user: {}\n",
46                    host, node.target_host, node.target_user
47                ));
48            }
49        }
50    }
51
52    yaml
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58    use consortium_nix::config::{AnsibleFleetConfig, DeploymentNode, ProfileType};
59
60    fn make_fleet(
61        nodes: Vec<(&str, &str, &str, Vec<&str>)>,
62        ansible: Option<AnsibleFleetConfig>,
63    ) -> FleetConfig {
64        let mut node_map = HashMap::new();
65        for (name, host, user, tags) in nodes {
66            node_map.insert(
67                name.to_string(),
68                DeploymentNode {
69                    name: name.to_string(),
70                    target_host: host.to_string(),
71                    target_user: user.to_string(),
72                    target_port: None,
73                    system: "x86_64-linux".to_string(),
74                    profile_type: ProfileType::Nixos,
75                    build_on_target: false,
76                    tags: tags.into_iter().map(|t| t.to_string()).collect(),
77                    drv_path: None,
78                    toplevel: None,
79                },
80            );
81        }
82        FleetConfig {
83            nodes: node_map,
84            builders: HashMap::new(),
85            flake_uri: ".".to_string(),
86            ansible_config: ansible,
87            slurm_config: None,
88            ray_config: None,
89            skypilot_config: None,
90        }
91    }
92
93    #[test]
94    fn test_generate_inventory_basic() {
95        let config = make_fleet(
96            vec![
97                ("web1", "10.0.0.1", "root", vec!["web"]),
98                ("db1", "10.0.0.2", "admin", vec!["db"]),
99            ],
100            None,
101        );
102        let inv = generate_inventory(&config);
103        assert!(inv.contains("web:"));
104        assert!(inv.contains("db:"));
105        assert!(inv.contains("ansible_host: 10.0.0.1"));
106        assert!(inv.contains("ansible_user: admin"));
107    }
108
109    #[test]
110    fn test_generate_inventory_with_host_groups() {
111        let config = make_fleet(
112            vec![
113                ("n1", "10.0.0.1", "root", vec![]),
114                ("n2", "10.0.0.2", "root", vec![]),
115            ],
116            Some(AnsibleFleetConfig {
117                control_node: "n1".to_string(),
118                ansible_version: None,
119                collections: vec![],
120                playbook_dir: None,
121                host_groups: {
122                    let mut g = HashMap::new();
123                    g.insert(
124                        "custom".to_string(),
125                        vec!["n1".to_string(), "n2".to_string()],
126                    );
127                    g
128                },
129            }),
130        );
131        let inv = generate_inventory(&config);
132        assert!(inv.contains("custom:"));
133    }
134
135    #[test]
136    fn test_generate_inventory_empty_fleet() {
137        let config = make_fleet(vec![], None);
138        let inv = generate_inventory(&config);
139        assert!(inv.contains("all:"));
140        assert!(inv.contains("children:"));
141    }
142
143    #[test]
144    fn test_generate_inventory_multi_tag_host() {
145        let config = make_fleet(
146            vec![("gpu1", "10.0.0.1", "root", vec!["gpu", "compute", "hpc"])],
147            None,
148        );
149        let inv = generate_inventory(&config);
150        assert!(inv.contains("gpu:"));
151        assert!(inv.contains("compute:"));
152        assert!(inv.contains("hpc:"));
153        // Host should appear under each tag group
154        let gpu1_count = inv.matches("gpu1:").count();
155        assert!(
156            gpu1_count >= 3,
157            "gpu1 should appear in 3 groups, found {}",
158            gpu1_count
159        );
160    }
161}
consortium_ansible/inventory.rs

consortium_ansible/
inventory.rs
