zombienet_provider/kubernetes/
provider.rs

1use std::{
2    collections::HashMap,
3    path::{Path, PathBuf},
4    sync::{Arc, Weak},
5};
6
7use async_trait::async_trait;
8use support::fs::FileSystem;
9use tokio::sync::RwLock;
10
11use super::{client::KubernetesClient, namespace::KubernetesNamespace};
12use crate::{
13    types::ProviderCapabilities, DynNamespace, Provider, ProviderError, ProviderNamespace,
14};
15
16const PROVIDER_NAME: &str = "k8s";
17
18pub struct KubernetesProvider<FS>
19where
20    FS: FileSystem + Send + Sync + Clone,
21{
22    weak: Weak<KubernetesProvider<FS>>,
23    capabilities: ProviderCapabilities,
24    tmp_dir: PathBuf,
25    k8s_client: KubernetesClient,
26    filesystem: FS,
27    pub(super) namespaces: RwLock<HashMap<String, Arc<KubernetesNamespace<FS>>>>,
28}
29
30impl<FS> KubernetesProvider<FS>
31where
32    FS: FileSystem + Send + Sync + Clone,
33{
34    pub async fn new(filesystem: FS) -> Arc<Self> {
35        let k8s_client = KubernetesClient::new().await.unwrap();
36
37        Arc::new_cyclic(|weak| KubernetesProvider {
38            weak: weak.clone(),
39            capabilities: ProviderCapabilities {
40                requires_image: true,
41                has_resources: true,
42                prefix_with_full_path: false,
43                use_default_ports_in_cmd: true,
44            },
45            tmp_dir: std::env::temp_dir(),
46            k8s_client,
47            filesystem,
48            namespaces: RwLock::new(HashMap::new()),
49        })
50    }
51
52    pub fn tmp_dir(mut self, tmp_dir: impl Into<PathBuf>) -> Self {
53        self.tmp_dir = tmp_dir.into();
54        self
55    }
56}
57
58#[async_trait]
59impl<FS> Provider for KubernetesProvider<FS>
60where
61    FS: FileSystem + Send + Sync + Clone + 'static,
62{
63    fn name(&self) -> &str {
64        PROVIDER_NAME
65    }
66
67    fn capabilities(&self) -> &ProviderCapabilities {
68        &self.capabilities
69    }
70
71    async fn namespaces(&self) -> HashMap<String, DynNamespace> {
72        self.namespaces
73            .read()
74            .await
75            .iter()
76            .map(|(name, namespace)| (name.clone(), namespace.clone() as DynNamespace))
77            .collect()
78    }
79
80    async fn create_namespace(&self) -> Result<DynNamespace, ProviderError> {
81        let namespace = KubernetesNamespace::new(
82            &self.weak,
83            &self.tmp_dir,
84            &self.capabilities,
85            &self.k8s_client,
86            &self.filesystem,
87            None,
88        )
89        .await?;
90
91        self.namespaces
92            .write()
93            .await
94            .insert(namespace.name().to_string(), namespace.clone());
95
96        Ok(namespace)
97    }
98
99    async fn create_namespace_with_base_dir(
100        &self,
101        base_dir: &Path,
102    ) -> Result<DynNamespace, ProviderError> {
103        let namespace = KubernetesNamespace::new(
104            &self.weak,
105            &self.tmp_dir,
106            &self.capabilities,
107            &self.k8s_client,
108            &self.filesystem,
109            Some(base_dir),
110        )
111        .await?;
112
113        self.namespaces
114            .write()
115            .await
116            .insert(namespace.name().to_string(), namespace.clone());
117
118        Ok(namespace)
119    }
120}