zombienet_sdk/
environment.rs

1//! Helpers functions to get configuration (e.g. Provider and images) from the env vars
2use std::{env, future::Future, path::PathBuf, pin::Pin};
3
4use crate::{
5    AttachToLive, AttachToLiveNetwork, LocalFileSystem, Network, NetworkConfig, NetworkConfigExt,
6    OrchestratorError,
7};
8
9const DEFAULT_POLKADOT_IMAGE: &str = "docker.io/parity/polkadot:latest";
10const DEFAULT_CUMULUS_IMAGE: &str = "docker.io/parity/polkadot-parachain:latest";
11
12#[derive(Debug, Default)]
13pub struct Images {
14    pub polkadot: String,
15    pub cumulus: String,
16}
17
18pub enum Provider {
19    Native,
20    K8s,
21    Docker,
22}
23
24impl Provider {
25    pub fn get_spawn_fn(
26        &self,
27    ) -> fn(NetworkConfig) -> Pin<Box<dyn Future<Output = SpawnResult> + Send>> {
28        match self {
29            Provider::Native => NetworkConfigExt::spawn_native,
30            Provider::K8s => NetworkConfigExt::spawn_k8s,
31            Provider::Docker => NetworkConfigExt::spawn_docker,
32        }
33    }
34}
35
36// Use `docker` as default provider
37impl From<String> for Provider {
38    fn from(value: String) -> Self {
39        match value.to_ascii_lowercase().as_ref() {
40            "native" => Provider::Native,
41            "k8s" => Provider::K8s,
42            _ => Provider::Docker, // default provider
43        }
44    }
45}
46
47pub fn get_images_from_env() -> Images {
48    let polkadot = env::var("POLKADOT_IMAGE").unwrap_or(DEFAULT_POLKADOT_IMAGE.into());
49    let cumulus = env::var("CUMULUS_IMAGE").unwrap_or(DEFAULT_CUMULUS_IMAGE.into());
50    Images { polkadot, cumulus }
51}
52
53pub fn get_provider_from_env() -> Provider {
54    env::var("ZOMBIE_PROVIDER").unwrap_or_default().into()
55}
56
57pub type SpawnResult = Result<Network<LocalFileSystem>, OrchestratorError>;
58pub fn get_spawn_fn() -> fn(NetworkConfig) -> Pin<Box<dyn Future<Output = SpawnResult> + Send>> {
59    let provider = get_provider_from_env();
60
61    match provider {
62        Provider::Native => NetworkConfigExt::spawn_native,
63        Provider::K8s => NetworkConfigExt::spawn_k8s,
64        Provider::Docker => NetworkConfigExt::spawn_docker,
65    }
66}
67
68pub type AttachResult = Result<Network<LocalFileSystem>, OrchestratorError>;
69
70pub fn get_attach_fn() -> fn(PathBuf) -> Pin<Box<dyn Future<Output = AttachResult> + Send>> {
71    let provider = get_provider_from_env();
72
73    match provider {
74        Provider::Native => AttachToLiveNetwork::attach_native,
75        Provider::K8s => AttachToLiveNetwork::attach_k8s,
76        Provider::Docker => AttachToLiveNetwork::attach_docker,
77    }
78}
zombienet_sdk/environment.rs

zombienet_sdk/
environment.rs
