zombienet_sdk/
environment.rs

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