zombienet_provider/shared/
types.rs

1use std::{
2    collections::HashMap,
3    path::{Path, PathBuf},
4    process::ExitStatus,
5};
6
7use configuration::{shared::resources::Resources, types::AssetLocation};
8use serde::{Deserialize, Serialize};
9
10pub type Port = u16;
11
12pub type ExecutionResult = Result<String, (ExitStatus, String)>;
13
14#[derive(Debug, Clone, PartialEq)]
15pub struct ProviderCapabilities {
16    // default ports internal
17    /// Ensure that we have an image for each node (k8s/podman/docker)
18    pub requires_image: bool,
19    /// Allow to customize the resources through manifest (k8s).
20    pub has_resources: bool,
21    /// Used in native to prefix filepath with fullpath
22    pub prefix_with_full_path: bool,
23    /// Use default ports in node cmd/args.
24    /// NOTE: generally used in k8s/dockers since the images expose those ports.
25    pub use_default_ports_in_cmd: bool,
26}
27
28#[derive(Debug, Clone)]
29pub struct SpawnNodeOptions {
30    /// Name of the node
31    pub name: String,
32    /// Image of the node (IFF is supported by the provider)
33    pub image: Option<String>,
34    /// Resources to apply to the node (IFF is supported by the provider)
35    pub resources: Option<Resources>,
36    /// Main command to execute
37    pub program: String,
38    /// Arguments to pass to the main command
39    pub args: Vec<String>,
40    /// Environment to set when running the `program`
41    pub env: Vec<(String, String)>,
42    // TODO: rename startup_files
43    /// Files to inject at startup
44    pub injected_files: Vec<TransferedFile>,
45    /// Paths to create before start the node (e.g keystore)
46    /// should be created with `create_dir_all` in order
47    /// to create the full path even when we have missing parts
48    pub created_paths: Vec<PathBuf>,
49    /// Database snapshot to be injected (should be a tgz file)
50    /// Could be a local or remote asset
51    pub db_snapshot: Option<AssetLocation>,
52    pub port_mapping: Option<HashMap<Port, Port>>,
53    /// Optionally specify a log path for the node
54    pub node_log_path: Option<PathBuf>,
55}
56
57impl SpawnNodeOptions {
58    pub fn new<S>(name: S, program: S) -> Self
59    where
60        S: AsRef<str>,
61    {
62        Self {
63            name: name.as_ref().to_string(),
64            image: None,
65            resources: None,
66            program: program.as_ref().to_string(),
67            args: vec![],
68            env: vec![],
69            injected_files: vec![],
70            created_paths: vec![],
71            db_snapshot: None,
72            port_mapping: None,
73            node_log_path: None,
74        }
75    }
76
77    pub fn image<S>(mut self, image: S) -> Self
78    where
79        S: AsRef<str>,
80    {
81        self.image = Some(image.as_ref().to_string());
82        self
83    }
84
85    pub fn resources(mut self, resources: Resources) -> Self {
86        self.resources = Some(resources);
87        self
88    }
89
90    pub fn db_snapshot(mut self, db_snap: Option<AssetLocation>) -> Self {
91        self.db_snapshot = db_snap;
92        self
93    }
94
95    pub fn args<S, I>(mut self, args: I) -> Self
96    where
97        S: AsRef<str>,
98        I: IntoIterator<Item = S>,
99    {
100        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
101        self
102    }
103
104    pub fn env<S, I>(mut self, env: I) -> Self
105    where
106        S: AsRef<str>,
107        I: IntoIterator<Item = (S, S)>,
108    {
109        self.env = env
110            .into_iter()
111            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
112            .collect();
113        self
114    }
115
116    pub fn injected_files<I>(mut self, injected_files: I) -> Self
117    where
118        I: IntoIterator<Item = TransferedFile>,
119    {
120        self.injected_files = injected_files.into_iter().collect();
121        self
122    }
123
124    pub fn created_paths<P, I>(mut self, created_paths: I) -> Self
125    where
126        P: AsRef<Path>,
127        I: IntoIterator<Item = P>,
128    {
129        self.created_paths = created_paths
130            .into_iter()
131            .map(|path| path.as_ref().into())
132            .collect();
133        self
134    }
135
136    pub fn port_mapping(mut self, ports: HashMap<Port, Port>) -> Self {
137        self.port_mapping = Some(ports);
138        self
139    }
140
141    pub fn node_log_path(mut self, path: Option<PathBuf>) -> Self {
142        self.node_log_path = path;
143        self
144    }
145}
146
147#[derive(Debug)]
148pub struct GenerateFileCommand {
149    pub program: String,
150    pub args: Vec<String>,
151    pub env: Vec<(String, String)>,
152    pub local_output_path: PathBuf,
153}
154
155impl GenerateFileCommand {
156    pub fn new<S, P>(program: S, local_output_path: P) -> Self
157    where
158        S: AsRef<str>,
159        P: AsRef<Path>,
160    {
161        Self {
162            program: program.as_ref().to_string(),
163            args: vec![],
164            env: vec![],
165            local_output_path: local_output_path.as_ref().into(),
166        }
167    }
168
169    pub fn args<S, I>(mut self, args: I) -> Self
170    where
171        S: AsRef<str>,
172        I: IntoIterator<Item = S>,
173    {
174        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
175        self
176    }
177
178    pub fn env<S, I>(mut self, env: I) -> Self
179    where
180        S: AsRef<str>,
181        I: IntoIterator<Item = (S, S)>,
182    {
183        self.env = env
184            .into_iter()
185            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
186            .collect();
187        self
188    }
189}
190
191#[derive(Debug)]
192pub struct GenerateFilesOptions {
193    pub commands: Vec<GenerateFileCommand>,
194    pub image: Option<String>,
195    pub injected_files: Vec<TransferedFile>,
196    // Allow to control the name of the node used to create the files.
197    pub temp_name: Option<String>,
198    pub expected_path: Option<PathBuf>,
199}
200
201impl GenerateFilesOptions {
202    pub fn new<I>(commands: I, image: Option<String>, expected_path: Option<PathBuf>) -> Self
203    where
204        I: IntoIterator<Item = GenerateFileCommand>,
205    {
206        Self {
207            commands: commands.into_iter().collect(),
208            injected_files: vec![],
209            image,
210            temp_name: None,
211            expected_path,
212        }
213    }
214
215    pub fn with_files<I>(
216        commands: I,
217        image: Option<String>,
218        injected_files: &[TransferedFile],
219        expected_path: Option<PathBuf>,
220    ) -> Self
221    where
222        I: IntoIterator<Item = GenerateFileCommand>,
223    {
224        Self {
225            commands: commands.into_iter().collect(),
226            injected_files: injected_files.into(),
227            image,
228            temp_name: None,
229            expected_path,
230        }
231    }
232
233    pub fn image<S>(mut self, image: S) -> Self
234    where
235        S: AsRef<str>,
236    {
237        self.image = Some(image.as_ref().to_string());
238        self
239    }
240
241    pub fn injected_files<I>(mut self, injected_files: I) -> Self
242    where
243        I: IntoIterator<Item = TransferedFile>,
244    {
245        self.injected_files = injected_files.into_iter().collect();
246        self
247    }
248
249    pub fn temp_name(mut self, name: impl Into<String>) -> Self {
250        self.temp_name = Some(name.into());
251        self
252    }
253}
254
255#[derive(Debug)]
256pub struct RunCommandOptions {
257    pub program: String,
258    pub args: Vec<String>,
259    pub env: Vec<(String, String)>,
260}
261
262impl RunCommandOptions {
263    pub fn new<S>(program: S) -> Self
264    where
265        S: AsRef<str>,
266    {
267        Self {
268            program: program.as_ref().to_string(),
269            args: vec![],
270            env: vec![],
271        }
272    }
273
274    pub fn args<S, I>(mut self, args: I) -> Self
275    where
276        S: AsRef<str>,
277        I: IntoIterator<Item = S>,
278    {
279        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
280        self
281    }
282
283    pub fn env<S, I>(mut self, env: I) -> Self
284    where
285        S: AsRef<str>,
286        I: IntoIterator<Item = (S, S)>,
287    {
288        self.env = env
289            .into_iter()
290            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
291            .collect();
292        self
293    }
294}
295
296pub struct RunScriptOptions {
297    pub local_script_path: PathBuf,
298    pub args: Vec<String>,
299    pub env: Vec<(String, String)>,
300}
301
302impl RunScriptOptions {
303    pub fn new<P>(local_script_path: P) -> Self
304    where
305        P: AsRef<Path>,
306    {
307        Self {
308            local_script_path: local_script_path.as_ref().into(),
309            args: vec![],
310            env: vec![],
311        }
312    }
313
314    pub fn args<S, I>(mut self, args: I) -> Self
315    where
316        S: AsRef<str>,
317        I: IntoIterator<Item = S>,
318    {
319        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
320        self
321    }
322
323    pub fn env<S, I>(mut self, env: I) -> Self
324    where
325        S: AsRef<str>,
326        I: IntoIterator<Item = (S, S)>,
327    {
328        self.env = env
329            .into_iter()
330            .map(|(name, value)| (name.as_ref().to_string(), value.as_ref().to_string()))
331            .collect();
332        self
333    }
334}
335
336// TODO(team): I think we can rename it to FileMap?
337#[derive(Debug, Clone, Serialize, Deserialize)]
338pub struct TransferedFile {
339    pub local_path: PathBuf,
340    pub remote_path: PathBuf,
341    // TODO: Can be narrowed to have strict typing on this?
342    pub mode: String,
343}
344
345impl TransferedFile {
346    pub fn new<P>(local_path: P, remote_path: P) -> Self
347    where
348        P: AsRef<Path>,
349    {
350        Self {
351            local_path: local_path.as_ref().into(),
352            remote_path: remote_path.as_ref().into(),
353            mode: "0644".to_string(), // default to rw-r--r--
354        }
355    }
356
357    pub fn mode<S>(mut self, mode: S) -> Self
358    where
359        S: AsRef<str>,
360    {
361        self.mode = mode.as_ref().to_string();
362        self
363    }
364}
365
366impl std::fmt::Display for TransferedFile {
367    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
368        write!(
369            f,
370            "File to transfer (local: {}, remote: {})",
371            self.local_path.display(),
372            self.remote_path.display()
373        )
374    }
375}
zombienet_provider/shared/types.rs

zombienet_provider/shared/
types.rs
