Skip to main content

zombienet_support/
replacer.rs

1use std::collections::{HashMap, HashSet};
2
3use lazy_static::lazy_static;
4use regex::{Captures, Regex};
5use tracing::{trace, warn};
6
7use crate::constants::{SHOULD_COMPILE, THIS_IS_A_BUG};
8
9lazy_static! {
10    static ref RE: Regex = Regex::new(r#"\{\{([a-zA-Z0-9_]*)\}\}"#)
11        .unwrap_or_else(|_| panic!("{SHOULD_COMPILE}, {THIS_IS_A_BUG}"));
12    static ref TOKEN_PLACEHOLDER: Regex = Regex::new(r#"\{\{ZOMBIE:(.*?):(.*?)\}\}"#)
13        .unwrap_or_else(|_| panic!("{SHOULD_COMPILE}, {THIS_IS_A_BUG}"));
14    static ref PLACEHOLDER_COMPAT: HashMap<&'static str, &'static str> = {
15        let mut m = HashMap::new();
16        m.insert("multiAddress", "multiaddr");
17        m.insert("wsUri", "ws_uri");
18        m.insert("prometheusUri", "prometheus_uri");
19
20        m
21    };
22}
23
24/// Return true if the text contains any TOKEN_PLACEHOLDER
25pub fn has_tokens(text: &str) -> bool {
26    TOKEN_PLACEHOLDER.is_match(text)
27}
28
29pub fn apply_replacements(text: &str, replacements: &HashMap<&str, &str>) -> String {
30    let augmented_text = RE.replace_all(text, |caps: &Captures| {
31        if let Some(replacements_value) = replacements.get(&caps[1]) {
32            replacements_value.to_string()
33        } else {
34            caps[0].to_string()
35        }
36    });
37
38    augmented_text.to_string()
39}
40
41pub fn apply_env_replacements(text: &str) -> String {
42    let augmented_text = RE.replace_all(text, |caps: &Captures| {
43        if let Ok(replacements_value) = std::env::var(&caps[1]) {
44            replacements_value
45        } else {
46            caps[0].to_string()
47        }
48    });
49
50    augmented_text.to_string()
51}
52
53pub fn apply_running_network_replacements(text: &str, network: &serde_json::Value) -> String {
54    let augmented_text = TOKEN_PLACEHOLDER.replace_all(text, |caps: &Captures| {
55        trace!("appling replacements for caps: {caps:#?}");
56        if let Some(node) = network.get(&caps[1]) {
57            trace!("caps1 {} - node: {node}", &caps[1]);
58            let field = *PLACEHOLDER_COMPAT.get(&caps[2]).unwrap_or(&&caps[2]);
59            if let Some(val) = node.get(field) {
60                trace!("caps2 {} - node: {node}", field);
61                val.as_str().unwrap_or("Invalid string").to_string()
62            } else {
63                warn!(
64                    "⚠️ The node with name {} doesn't have the value {} in context",
65                    &caps[1], &caps[2]
66                );
67                // DEBUG
68                trace!("node: {:?}", node);
69                caps[0].to_string()
70            }
71        } else {
72            warn!("⚠️ No node with name {} in context", &caps[1]);
73            caps[0].to_string()
74        }
75    });
76
77    augmented_text.to_string()
78}
79
80pub fn get_tokens_to_replace(text: &str) -> HashSet<String> {
81    let mut tokens = HashSet::new();
82
83    TOKEN_PLACEHOLDER
84        .captures_iter(text)
85        .for_each(|caps: Captures| {
86            tokens.insert(caps[1].to_string());
87        });
88
89    tokens
90}
91
92#[cfg(test)]
93mod tests {
94    use serde_json::json;
95
96    use super::*;
97
98    #[test]
99    fn replace_should_works() {
100        let text = "some {{namespace}}";
101        let mut replacements = HashMap::new();
102        replacements.insert("namespace", "demo-123");
103        let res = apply_replacements(text, &replacements);
104        assert_eq!("some demo-123".to_string(), res);
105    }
106
107    #[test]
108    fn replace_env_should_works() {
109        let text = "some {{namespace}}";
110        std::env::set_var("namespace", "demo-123");
111        // let mut replacements = HashMap::new();
112        // replacements.insert("namespace", "demo-123");
113        let res = apply_env_replacements(text);
114        assert_eq!("some demo-123".to_string(), res);
115    }
116
117    #[test]
118    fn replace_multiple_should_works() {
119        let text = r#"some {{namespace}}
120        other is {{other}}"#;
121        let augmented_text = r#"some demo-123
122        other is other-123"#;
123
124        let mut replacements = HashMap::new();
125        replacements.insert("namespace", "demo-123");
126        replacements.insert("other", "other-123");
127        let res = apply_replacements(text, &replacements);
128        assert_eq!(augmented_text, res);
129    }
130
131    #[test]
132    fn replace_multiple_with_missing_should_works() {
133        let text = r#"some {{namespace}}
134        other is {{other}}"#;
135        let augmented_text = r#"some demo-123
136        other is {{other}}"#;
137
138        let mut replacements = HashMap::new();
139        replacements.insert("namespace", "demo-123");
140
141        let res = apply_replacements(text, &replacements);
142        assert_eq!(augmented_text, res);
143    }
144
145    #[test]
146    fn replace_without_replacement_should_leave_text_unchanged() {
147        let text = "some {{namespace}}";
148        let mut replacements = HashMap::new();
149        replacements.insert("other", "demo-123");
150        let res = apply_replacements(text, &replacements);
151        assert_eq!(text.to_string(), res);
152    }
153
154    #[test]
155    fn replace_running_network_should_work() {
156        let network = json!({
157            "alice" : {
158                "multiaddr": "some/demo/127.0.0.1"
159            }
160        });
161
162        let res = apply_running_network_replacements("ws://{{ZOMBIE:alice:multiaddr}}", &network);
163        assert_eq!(res.as_str(), "ws://some/demo/127.0.0.1");
164    }
165
166    #[test]
167    fn replace_running_network_with_compat_should_work() {
168        let network = json!({
169            "alice" : {
170                "multiaddr": "some/demo/127.0.0.1"
171            }
172        });
173
174        let res = apply_running_network_replacements("{{ZOMBIE:alice:multiAddress}}", &network);
175        assert_eq!(res.as_str(), "some/demo/127.0.0.1");
176    }
177
178    #[test]
179    fn replace_running_network_with_missing_field_should_not_replace_nothing() {
180        let network = json!({
181            "alice" : {
182                "multiaddr": "some/demo/127.0.0.1"
183            }
184        });
185
186        let res = apply_running_network_replacements("{{ZOMBIE:alice:someField}}", &network);
187        assert_eq!(res.as_str(), "{{ZOMBIE:alice:someField}}");
188    }
189
190    #[test]
191    fn get_tokens_to_replace_should_work() {
192        let res = get_tokens_to_replace("{{ZOMBIE:alice:multiaddr}} {{ZOMBIE:bob:multiaddr}}");
193        let mut expected = HashSet::new();
194        expected.insert("alice".to_string());
195        expected.insert("bob".to_string());
196
197        assert_eq!(res, expected);
198    }
199}