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";
11const DEFAULT_MALUS_IMAGE: &str = "docker.io/paritypr/malus:latest";
12const DEFAULT_COLANDER_IMAGE: &str = "docker.io/paritypr/colander:latest";
13
14#[derive(Debug, Default)]
15pub struct Images {
16    pub polkadot: String,
17    pub cumulus: String,
18    pub malus: String,
19    pub colander: String,
20}
21
22pub enum Provider {
23    Native,
24    K8s,
25    Docker,
26}
27
28impl Provider {
29    pub fn get_spawn_fn(
30        &self,
31    ) -> fn(NetworkConfig) -> Pin<Box<dyn Future<Output = SpawnResult> + Send>> {
32        match self {
33            Provider::Native => NetworkConfigExt::spawn_native,
34            Provider::K8s => NetworkConfigExt::spawn_k8s,
35            Provider::Docker => NetworkConfigExt::spawn_docker,
36        }
37    }
38}
39
40// Use `docker` as default provider
41impl From<String> for Provider {
42    fn from(value: String) -> Self {
43        match value.to_ascii_lowercase().as_ref() {
44            "native" => Provider::Native,
45            "k8s" => Provider::K8s,
46            _ => Provider::Docker, // default provider
47        }
48    }
49}
50
51pub fn get_images_from_env() -> Images {
52    let polkadot = env::var("POLKADOT_IMAGE").unwrap_or(DEFAULT_POLKADOT_IMAGE.into());
53    let cumulus = env::var("CUMULUS_IMAGE").unwrap_or(DEFAULT_CUMULUS_IMAGE.into());
54    let malus = env::var("MALUS_IMAGE").unwrap_or(DEFAULT_MALUS_IMAGE.into());
55    let colander = env::var("COL_IMAGE").unwrap_or(DEFAULT_COLANDER_IMAGE.into());
56
57    Images {
58        polkadot,
59        cumulus,
60        malus,
61        colander,
62    }
63}
64
65pub fn get_provider_from_env() -> Provider {
66    env::var("ZOMBIE_PROVIDER").unwrap_or_default().into()
67}
68
69pub type SpawnResult = Result<Network<LocalFileSystem>, OrchestratorError>;
70pub fn get_spawn_fn() -> fn(NetworkConfig) -> Pin<Box<dyn Future<Output = SpawnResult> + Send>> {
71    let provider = get_provider_from_env();
72
73    match provider {
74        Provider::Native => NetworkConfigExt::spawn_native,
75        Provider::K8s => NetworkConfigExt::spawn_k8s,
76        Provider::Docker => NetworkConfigExt::spawn_docker,
77    }
78}
79
80pub type AttachResult = Result<Network<LocalFileSystem>, OrchestratorError>;
81
82pub fn get_attach_fn() -> fn(PathBuf) -> Pin<Box<dyn Future<Output = AttachResult> + Send>> {
83    let provider = get_provider_from_env();
84
85    match provider {
86        Provider::Native => AttachToLiveNetwork::attach_native,
87        Provider::K8s => AttachToLiveNetwork::attach_k8s,
88        Provider::Docker => AttachToLiveNetwork::attach_docker,
89    }
90}