zombienet_orchestrator/
network.rs

1pub mod chain_upgrade;
2pub mod node;
3pub mod parachain;
4pub mod relaychain;
5
6use std::{cell::RefCell, collections::HashMap, path::PathBuf, rc::Rc};
7
8use configuration::{
9    para_states::{Initial, Running},
10    shared::{helpers::generate_unique_node_name_from_names, node::EnvVar},
11    types::{Arg, Command, Image, Port, ValidationContext},
12    ParachainConfig, ParachainConfigBuilder, RegistrationStrategy,
13};
14use provider::{types::TransferedFile, DynNamespace, ProviderError};
15use serde::Serialize;
16use support::fs::FileSystem;
17
18use self::{node::NetworkNode, parachain::Parachain, relaychain::Relaychain};
19use crate::{
20    generators::chain_spec::ChainSpec,
21    network_spec::{self, NetworkSpec},
22    shared::{
23        macros,
24        types::{ChainDefaultContext, RegisterParachainOptions},
25    },
26    spawner::{self, SpawnNodeCtx},
27    ScopedFilesystem, ZombieRole,
28};
29
30#[derive(Serialize)]
31pub struct Network<T: FileSystem> {
32    #[serde(skip)]
33    ns: DynNamespace,
34    #[serde(skip)]
35    filesystem: T,
36    relay: Relaychain,
37    initial_spec: NetworkSpec,
38    parachains: HashMap<u32, Vec<Parachain>>,
39    #[serde(skip)]
40    nodes_by_name: HashMap<String, NetworkNode>,
41}
42
43impl<T: FileSystem> std::fmt::Debug for Network<T> {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        f.debug_struct("Network")
46            .field("ns", &"ns_skipped")
47            .field("relay", &self.relay)
48            .field("initial_spec", &self.initial_spec)
49            .field("parachains", &self.parachains)
50            .field("nodes_by_name", &self.nodes_by_name)
51            .finish()
52    }
53}
54
55macros::create_add_options!(AddNodeOptions {
56    chain_spec: Option<PathBuf>
57});
58
59macros::create_add_options!(AddCollatorOptions {
60    chain_spec: Option<PathBuf>,
61    chain_spec_relay: Option<PathBuf>
62});
63
64impl<T: FileSystem> Network<T> {
65    pub(crate) fn new_with_relay(
66        relay: Relaychain,
67        ns: DynNamespace,
68        fs: T,
69        initial_spec: NetworkSpec,
70    ) -> Self {
71        Self {
72            ns,
73            filesystem: fs,
74            relay,
75            initial_spec,
76            parachains: Default::default(),
77            nodes_by_name: Default::default(),
78        }
79    }
80
81    // Pubic API
82    pub fn ns_name(&self) -> String {
83        self.ns.name().to_string()
84    }
85
86    pub fn base_dir(&self) -> Option<&str> {
87        self.ns.base_dir().to_str()
88    }
89
90    pub fn relaychain(&self) -> &Relaychain {
91        &self.relay
92    }
93
94    // Teardown the network
95    pub async fn destroy(self) -> Result<(), ProviderError> {
96        self.ns.destroy().await
97    }
98
99    /// Add a node to the relaychain
100    // The new node is added to the running network instance.
101    /// # Example:
102    /// ```rust
103    /// # use provider::NativeProvider;
104    /// # use support::{fs::local::LocalFileSystem};
105    /// # use zombienet_orchestrator::{errors, AddNodeOptions, Orchestrator};
106    /// # use configuration::NetworkConfig;
107    /// # async fn example() -> Result<(), errors::OrchestratorError> {
108    /// #   let provider = NativeProvider::new(LocalFileSystem {});
109    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
110    /// #   let config = NetworkConfig::load_from_toml("config.toml")?;
111    /// let mut network = orchestrator.spawn(config).await?;
112    ///
113    /// // Create the options to add the new node
114    /// let opts = AddNodeOptions {
115    ///     rpc_port: Some(9444),
116    ///     is_validator: true,
117    ///     ..Default::default()
118    /// };
119    ///
120    /// network.add_node("new-node", opts).await?;
121    /// #   Ok(())
122    /// # }
123    /// ```
124    pub async fn add_node(
125        &mut self,
126        name: impl Into<String>,
127        options: AddNodeOptions,
128    ) -> Result<(), anyhow::Error> {
129        let name = generate_unique_node_name_from_names(
130            name,
131            &mut self.nodes_by_name.keys().cloned().collect(),
132        );
133
134        let relaychain = self.relaychain();
135
136        let chain_spec_path = if let Some(chain_spec_custom_path) = &options.chain_spec {
137            chain_spec_custom_path.clone()
138        } else {
139            PathBuf::from(format!(
140                "{}/{}.json",
141                self.ns.base_dir().to_string_lossy(),
142                relaychain.chain
143            ))
144        };
145
146        let chain_context = ChainDefaultContext {
147            default_command: self.initial_spec.relaychain.default_command.as_ref(),
148            default_image: self.initial_spec.relaychain.default_image.as_ref(),
149            default_resources: self.initial_spec.relaychain.default_resources.as_ref(),
150            default_db_snapshot: self.initial_spec.relaychain.default_db_snapshot.as_ref(),
151            default_args: self.initial_spec.relaychain.default_args.iter().collect(),
152        };
153
154        let mut node_spec = network_spec::node::NodeSpec::from_ad_hoc(
155            &name,
156            options.into(),
157            &chain_context,
158            false,
159        )?;
160
161        node_spec.available_args_output = Some(
162            self.initial_spec
163                .node_available_args_output(&node_spec, self.ns.clone())
164                .await?,
165        );
166
167        let base_dir = self.ns.base_dir().to_string_lossy();
168        let scoped_fs = ScopedFilesystem::new(&self.filesystem, &base_dir);
169
170        let ctx = SpawnNodeCtx {
171            chain_id: &relaychain.chain_id,
172            parachain_id: None,
173            chain: &relaychain.chain,
174            role: ZombieRole::Node,
175            ns: &self.ns,
176            scoped_fs: &scoped_fs,
177            parachain: None,
178            bootnodes_addr: &vec![],
179            wait_ready: true,
180            nodes_by_name: serde_json::to_value(&self.nodes_by_name)?,
181        };
182
183        let global_files_to_inject = vec![TransferedFile::new(
184            chain_spec_path,
185            PathBuf::from(format!("/cfg/{}.json", relaychain.chain)),
186        )];
187
188        let node = spawner::spawn_node(&node_spec, global_files_to_inject, &ctx).await?;
189
190        // TODO: register the new node as validator in the relaychain
191        // STEPS:
192        //  - check balance of `stash` derivation for validator account
193        //  - call rotate_keys on the new validator
194        //  - call setKeys on the new validator
195        // if node_spec.is_validator {
196        //     let running_node = self.relay.nodes.first().unwrap();
197        //     // tx_helper::validator_actions::register(vec![&node], &running_node.ws_uri, None).await?;
198        // }
199
200        // Add node to relaychain data
201        self.add_running_node(node.clone(), None);
202
203        Ok(())
204    }
205
206    /// Add a new collator to a parachain
207    ///
208    /// NOTE: if more parachains with given id available (rare corner case)
209    /// then it adds collator to the first parachain
210    ///
211    /// # Example:
212    /// ```rust
213    /// # use provider::NativeProvider;
214    /// # use support::{fs::local::LocalFileSystem};
215    /// # use zombienet_orchestrator::{errors, AddCollatorOptions, Orchestrator};
216    /// # use configuration::NetworkConfig;
217    /// # async fn example() -> Result<(), anyhow::Error> {
218    /// #   let provider = NativeProvider::new(LocalFileSystem {});
219    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
220    /// #   let config = NetworkConfig::load_from_toml("config.toml")?;
221    /// let mut network = orchestrator.spawn(config).await?;
222    ///
223    /// let col_opts = AddCollatorOptions {
224    ///     command: Some("polkadot-parachain".try_into()?),
225    ///     ..Default::default()
226    /// };
227    ///
228    /// network.add_collator("new-col-1", col_opts, 100).await?;
229    /// #   Ok(())
230    /// # }
231    /// ```
232    pub async fn add_collator(
233        &mut self,
234        name: impl Into<String>,
235        options: AddCollatorOptions,
236        para_id: u32,
237    ) -> Result<(), anyhow::Error> {
238        let name = generate_unique_node_name_from_names(
239            name,
240            &mut self.nodes_by_name.keys().cloned().collect(),
241        );
242        let spec = self
243            .initial_spec
244            .parachains
245            .iter()
246            .find(|para| para.id == para_id)
247            .ok_or(anyhow::anyhow!(format!("parachain: {para_id} not found!")))?;
248        let role = if spec.is_cumulus_based {
249            ZombieRole::CumulusCollator
250        } else {
251            ZombieRole::Collator
252        };
253        let chain_context = ChainDefaultContext {
254            default_command: spec.default_command.as_ref(),
255            default_image: spec.default_image.as_ref(),
256            default_resources: spec.default_resources.as_ref(),
257            default_db_snapshot: spec.default_db_snapshot.as_ref(),
258            default_args: spec.default_args.iter().collect(),
259        };
260
261        let parachain = self
262            .parachains
263            .get_mut(&para_id)
264            .ok_or(anyhow::anyhow!(format!("parachain: {para_id} not found!")))?
265            .get_mut(0)
266            .ok_or(anyhow::anyhow!(format!("parachain: {para_id} not found!")))?;
267
268        let base_dir = self.ns.base_dir().to_string_lossy();
269        let scoped_fs = ScopedFilesystem::new(&self.filesystem, &base_dir);
270
271        // TODO: we want to still supporting spawn a dedicated bootnode??
272        let ctx = SpawnNodeCtx {
273            chain_id: &self.relay.chain_id,
274            parachain_id: parachain.chain_id.as_deref(),
275            chain: &self.relay.chain,
276            role,
277            ns: &self.ns,
278            scoped_fs: &scoped_fs,
279            parachain: Some(spec),
280            bootnodes_addr: &vec![],
281            wait_ready: true,
282            nodes_by_name: serde_json::to_value(&self.nodes_by_name)?,
283        };
284
285        let relaychain_spec_path = if let Some(chain_spec_custom_path) = &options.chain_spec_relay {
286            chain_spec_custom_path.clone()
287        } else {
288            PathBuf::from(format!(
289                "{}/{}.json",
290                self.ns.base_dir().to_string_lossy(),
291                self.relay.chain
292            ))
293        };
294
295        let mut global_files_to_inject = vec![TransferedFile::new(
296            relaychain_spec_path,
297            PathBuf::from(format!("/cfg/{}.json", self.relay.chain)),
298        )];
299
300        let para_chain_spec_local_path = if let Some(para_chain_spec_custom) = &options.chain_spec {
301            Some(para_chain_spec_custom.clone())
302        } else if let Some(para_spec_path) = &parachain.chain_spec_path {
303            Some(PathBuf::from(format!(
304                "{}/{}",
305                self.ns.base_dir().to_string_lossy(),
306                para_spec_path.to_string_lossy()
307            )))
308        } else {
309            None
310        };
311
312        if let Some(para_spec_path) = para_chain_spec_local_path {
313            global_files_to_inject.push(TransferedFile::new(
314                para_spec_path,
315                PathBuf::from(format!("/cfg/{para_id}.json")),
316            ));
317        }
318
319        let mut node_spec =
320            network_spec::node::NodeSpec::from_ad_hoc(name, options.into(), &chain_context, true)?;
321
322        node_spec.available_args_output = Some(
323            self.initial_spec
324                .node_available_args_output(&node_spec, self.ns.clone())
325                .await?,
326        );
327
328        let node = spawner::spawn_node(&node_spec, global_files_to_inject, &ctx).await?;
329        parachain.collators.push(node.clone());
330        self.add_running_node(node, None);
331
332        Ok(())
333    }
334
335    /// Get a parachain config builder from a running network
336    ///
337    /// This allow you to build a new parachain config to be deployed into
338    /// the running network.
339    pub fn para_config_builder(&self) -> ParachainConfigBuilder<Initial, Running> {
340        let used_ports = vec![]; // TODO: generate used ports from the network
341        let used_nodes_names = self.nodes_by_name.keys().cloned().collect();
342        let used_para_ids = HashMap::new(); // TODO: generate used para ids from the network
343
344        let context = ValidationContext {
345            used_ports,
346            used_nodes_names,
347            used_para_ids,
348        };
349        let context = Rc::new(RefCell::new(context));
350
351        ParachainConfigBuilder::new_with_running(context)
352    }
353
354    /// Add a new parachain to the running network
355    ///
356    /// # Arguments
357    /// * `para_config` - Parachain configuration to deploy
358    /// * `custom_relaychain_spec` - Optional path to a custom relaychain spec to use
359    /// * `custom_parchain_fs_prefix` - Optional prefix to use when artifacts are created
360    ///
361    ///
362    /// # Example:
363    /// ```rust
364    /// # use anyhow::anyhow;
365    /// # use provider::NativeProvider;
366    /// # use support::{fs::local::LocalFileSystem};
367    /// # use zombienet_orchestrator::{errors, AddCollatorOptions, Orchestrator};
368    /// # use configuration::NetworkConfig;
369    /// # async fn example() -> Result<(), anyhow::Error> {
370    /// #   let provider = NativeProvider::new(LocalFileSystem {});
371    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
372    /// #   let config = NetworkConfig::load_from_toml("config.toml")?;
373    /// let mut network = orchestrator.spawn(config).await?;
374    /// let para_config = network
375    ///     .para_config_builder()
376    ///     .with_id(100)
377    ///     .with_default_command("polkadot-parachain")
378    ///     .with_collator(|c| c.with_name("col-100-1"))
379    ///     .build()
380    ///     .map_err(|_e| anyhow!("Building config"))?;
381    ///
382    /// network.add_parachain(&para_config, None, None).await?;
383    ///
384    /// #   Ok(())
385    /// # }
386    /// ```
387    pub async fn add_parachain(
388        &mut self,
389        para_config: &ParachainConfig,
390        custom_relaychain_spec: Option<PathBuf>,
391        custom_parchain_fs_prefix: Option<String>,
392    ) -> Result<(), anyhow::Error> {
393        // build
394        let mut para_spec = network_spec::parachain::ParachainSpec::from_config(para_config)?;
395        let base_dir = self.ns.base_dir().to_string_lossy().to_string();
396        let scoped_fs = ScopedFilesystem::new(&self.filesystem, &base_dir);
397
398        let mut global_files_to_inject = vec![];
399
400        // get relaychain id
401        let relay_chain_id = if let Some(custom_path) = custom_relaychain_spec {
402            // use this file as relaychain spec
403            global_files_to_inject.push(TransferedFile::new(
404                custom_path.clone(),
405                PathBuf::from(format!("/cfg/{}.json", self.relaychain().chain)),
406            ));
407            let content = std::fs::read_to_string(custom_path)?;
408            ChainSpec::chain_id_from_spec(&content)?
409        } else {
410            global_files_to_inject.push(TransferedFile::new(
411                PathBuf::from(format!(
412                    "{}/{}",
413                    scoped_fs.base_dir,
414                    self.relaychain().chain_spec_path.to_string_lossy()
415                )),
416                PathBuf::from(format!("/cfg/{}.json", self.relaychain().chain)),
417            ));
418            self.relay.chain_id.clone()
419        };
420
421        let chain_spec_raw_path = para_spec
422            .build_chain_spec(&relay_chain_id, &self.ns, &scoped_fs)
423            .await?;
424
425        // Para artifacts
426        let para_path_prefix = if let Some(custom_prefix) = custom_parchain_fs_prefix {
427            custom_prefix
428        } else {
429            para_spec.id.to_string()
430        };
431
432        scoped_fs.create_dir(&para_path_prefix).await?;
433        // create wasm/state
434        para_spec
435            .genesis_state
436            .build(
437                chain_spec_raw_path.as_ref(),
438                format!("{}/genesis-state", &para_path_prefix),
439                &self.ns,
440                &scoped_fs,
441            )
442            .await?;
443        para_spec
444            .genesis_wasm
445            .build(
446                chain_spec_raw_path.as_ref(),
447                format!("{}/para_spec-wasm", &para_path_prefix),
448                &self.ns,
449                &scoped_fs,
450            )
451            .await?;
452
453        let parachain =
454            Parachain::from_spec(&para_spec, &global_files_to_inject, &scoped_fs).await?;
455        let parachain_id = parachain.chain_id.clone();
456
457        // Create `ctx` for spawn the nodes
458        let ctx_para = SpawnNodeCtx {
459            parachain: Some(&para_spec),
460            parachain_id: parachain_id.as_deref(),
461            role: if para_spec.is_cumulus_based {
462                ZombieRole::CumulusCollator
463            } else {
464                ZombieRole::Collator
465            },
466            bootnodes_addr: &vec![],
467            chain_id: &self.relaychain().chain_id,
468            chain: &self.relaychain().chain,
469            ns: &self.ns,
470            scoped_fs: &scoped_fs,
471            wait_ready: false,
472            nodes_by_name: serde_json::to_value(&self.nodes_by_name)?,
473        };
474
475        // Register the parachain to the running network
476        let first_node_url = self
477            .relaychain()
478            .nodes
479            .first()
480            .ok_or(anyhow::anyhow!(
481                "At least one node of the relaychain should be running"
482            ))?
483            .ws_uri();
484
485        if para_config.registration_strategy() == Some(&RegistrationStrategy::UsingExtrinsic) {
486            let register_para_options = RegisterParachainOptions {
487                id: parachain.para_id,
488                // This needs to resolve correctly
489                wasm_path: para_spec
490                    .genesis_wasm
491                    .artifact_path()
492                    .ok_or(anyhow::anyhow!(
493                        "artifact path for wasm must be set at this point",
494                    ))?
495                    .to_path_buf(),
496                state_path: para_spec
497                    .genesis_state
498                    .artifact_path()
499                    .ok_or(anyhow::anyhow!(
500                        "artifact path for state must be set at this point",
501                    ))?
502                    .to_path_buf(),
503                node_ws_url: first_node_url.to_string(),
504                onboard_as_para: para_spec.onboard_as_parachain,
505                seed: None, // TODO: Seed is passed by?
506                finalization: false,
507            };
508
509            Parachain::register(register_para_options, &scoped_fs).await?;
510        }
511
512        // Spawn the nodes
513        let spawning_tasks = para_spec
514            .collators
515            .iter()
516            .map(|node| spawner::spawn_node(node, parachain.files_to_inject.clone(), &ctx_para));
517
518        let running_nodes = futures::future::try_join_all(spawning_tasks).await?;
519        let running_para_id = parachain.para_id;
520        self.add_para(parachain);
521        for node in running_nodes {
522            self.add_running_node(node, Some(running_para_id));
523        }
524
525        Ok(())
526    }
527
528    /// Register a parachain, which has already been added to the network (with manual registration
529    /// strategy)
530    ///
531    /// # Arguments
532    /// * `para_id` - Parachain Id
533    ///
534    ///
535    /// # Example:
536    /// ```rust
537    /// # use anyhow::anyhow;
538    /// # use provider::NativeProvider;
539    /// # use support::{fs::local::LocalFileSystem};
540    /// # use zombienet_orchestrator::Orchestrator;
541    /// # use configuration::{NetworkConfig, NetworkConfigBuilder, RegistrationStrategy};
542    /// # async fn example() -> Result<(), anyhow::Error> {
543    /// #   let provider = NativeProvider::new(LocalFileSystem {});
544    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
545    /// #   let config = NetworkConfigBuilder::new()
546    /// #     .with_relaychain(|r| {
547    /// #       r.with_chain("rococo-local")
548    /// #         .with_default_command("polkadot")
549    /// #         .with_node(|node| node.with_name("alice"))
550    /// #     })
551    /// #     .with_parachain(|p| {
552    /// #       p.with_id(100)
553    /// #         .with_registration_strategy(RegistrationStrategy::Manual)
554    /// #         .with_default_command("test-parachain")
555    /// #         .with_collator(|n| n.with_name("dave").validator(false))
556    /// #     })
557    /// #     .build()
558    /// #     .map_err(|_e| anyhow!("Building config"))?;
559    /// let mut network = orchestrator.spawn(config).await?;
560    ///
561    /// network.register_parachain(100).await?;
562    ///
563    /// #   Ok(())
564    /// # }
565    /// ```
566    pub async fn register_parachain(&mut self, para_id: u32) -> Result<(), anyhow::Error> {
567        let para = self
568            .initial_spec
569            .parachains
570            .iter()
571            .find(|p| p.id == para_id)
572            .ok_or(anyhow::anyhow!(
573                "no parachain with id = {para_id} available",
574            ))?;
575        let para_genesis_config = para.get_genesis_config()?;
576        let first_node_url = self
577            .relaychain()
578            .nodes
579            .first()
580            .ok_or(anyhow::anyhow!(
581                "At least one node of the relaychain should be running"
582            ))?
583            .ws_uri();
584        let register_para_options: RegisterParachainOptions = RegisterParachainOptions {
585            id: para_id,
586            // This needs to resolve correctly
587            wasm_path: para_genesis_config.wasm_path.clone(),
588            state_path: para_genesis_config.state_path.clone(),
589            node_ws_url: first_node_url.to_string(),
590            onboard_as_para: para_genesis_config.as_parachain,
591            seed: None, // TODO: Seed is passed by?
592            finalization: false,
593        };
594        let base_dir = self.ns.base_dir().to_string_lossy().to_string();
595        let scoped_fs = ScopedFilesystem::new(&self.filesystem, &base_dir);
596        Parachain::register(register_para_options, &scoped_fs).await?;
597
598        Ok(())
599    }
600
601    // deregister and stop the collator?
602    // remove_parachain()
603
604    pub fn get_node(&self, name: impl Into<String>) -> Result<&NetworkNode, anyhow::Error> {
605        let name = name.into();
606        if let Some(node) = self.nodes_iter().find(|&n| n.name == name) {
607            return Ok(node);
608        }
609
610        let list = self
611            .nodes_iter()
612            .map(|n| &n.name)
613            .cloned()
614            .collect::<Vec<_>>()
615            .join(", ");
616
617        Err(anyhow::anyhow!(
618            "can't find node with name: {name:?}, should be one of {list}"
619        ))
620    }
621
622    pub fn get_node_mut(
623        &mut self,
624        name: impl Into<String>,
625    ) -> Result<&mut NetworkNode, anyhow::Error> {
626        let name = name.into();
627        self.nodes_iter_mut()
628            .find(|n| n.name == name)
629            .ok_or(anyhow::anyhow!("can't find node with name: {name:?}"))
630    }
631
632    pub fn nodes(&self) -> Vec<&NetworkNode> {
633        self.nodes_by_name.values().collect::<Vec<&NetworkNode>>()
634    }
635
636    pub async fn detach(&self) {
637        self.ns.detach().await
638    }
639
640    // Internal API
641    pub(crate) fn add_running_node(&mut self, node: NetworkNode, para_id: Option<u32>) {
642        if let Some(para_id) = para_id {
643            if let Some(para) = self.parachains.get_mut(&para_id).and_then(|p| p.get_mut(0)) {
644                para.collators.push(node.clone());
645            } else {
646                // is the first node of the para, let create the entry
647                unreachable!()
648            }
649        } else {
650            self.relay.nodes.push(node.clone());
651        }
652        // TODO: we should hold a ref to the node in the vec in the future.
653        let node_name = node.name.clone();
654        self.nodes_by_name.insert(node_name, node);
655    }
656
657    pub(crate) fn add_para(&mut self, para: Parachain) {
658        self.parachains.entry(para.para_id).or_default().push(para);
659    }
660
661    pub fn name(&self) -> &str {
662        self.ns.name()
663    }
664
665    /// Get a first parachain from the list of the parachains with specified id.
666    /// NOTE!
667    /// Usually the list will contain only one parachain.
668    /// Multiple parachains with the same id is a corner case.
669    /// If this is the case then one can get such parachain with
670    /// `parachain_by_unique_id()` method
671    ///
672    /// # Arguments
673    /// * `para_id` - Parachain Id
674    pub fn parachain(&self, para_id: u32) -> Option<&Parachain> {
675        self.parachains.get(&para_id)?.first()
676    }
677
678    /// Get a parachain by its unique id.
679    ///
680    /// This is particularly useful if there are multiple parachains
681    /// with the same id (this is a rare corner case).
682    ///
683    /// # Arguments
684    /// * `unique_id` - unique id of the parachain
685    pub fn parachain_by_unique_id(&self, unique_id: impl AsRef<str>) -> Option<&Parachain> {
686        self.parachains
687            .values()
688            .flat_map(|p| p.iter())
689            .find(|p| p.unique_id == unique_id.as_ref())
690    }
691
692    pub fn parachains(&self) -> Vec<&Parachain> {
693        self.parachains.values().flatten().collect()
694    }
695
696    pub(crate) fn nodes_iter(&self) -> impl Iterator<Item = &NetworkNode> {
697        self.relay.nodes.iter().chain(
698            self.parachains
699                .values()
700                .flat_map(|p| p.iter())
701                .flat_map(|p| &p.collators),
702        )
703    }
704
705    pub(crate) fn nodes_iter_mut(&mut self) -> impl Iterator<Item = &mut NetworkNode> {
706        self.relay.nodes.iter_mut().chain(
707            self.parachains
708                .values_mut()
709                .flat_map(|p| p.iter_mut())
710                .flat_map(|p| &mut p.collators),
711        )
712    }
713
714    /// Waits given number of seconds until all nodes in the network report that they are
715    /// up and running.
716    ///
717    /// # Arguments
718    /// * `timeout_secs` - The number of seconds to wait.
719    ///
720    /// # Returns
721    /// * `Ok()` if the node is up before timeout occured.
722    /// * `Err(e)` if timeout or other error occurred while waiting.
723    pub async fn wait_until_is_up(&self, timeout_secs: u64) -> Result<(), anyhow::Error> {
724        let handles = self
725            .nodes_iter()
726            .map(|node| node.wait_until_is_up(timeout_secs));
727
728        futures::future::try_join_all(handles).await?;
729
730        Ok(())
731    }
732}
zombienet_orchestrator/network.rs

zombienet_orchestrator/
network.rs
