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    shared::helpers::extract_namespace_info, types::ProviderCapabilities, DynNamespace, Provider,
14    ProviderError, ProviderNamespace,
15};
16
17pub const PROVIDER_NAME: &str = "k8s";
18
19pub struct KubernetesProvider<FS>
20where
21    FS: FileSystem + Send + Sync + Clone,
22{
23    weak: Weak<KubernetesProvider<FS>>,
24    capabilities: ProviderCapabilities,
25    tmp_dir: PathBuf,
26    k8s_client: KubernetesClient,
27    filesystem: FS,
28    pub(super) namespaces: RwLock<HashMap<String, Arc<KubernetesNamespace<FS>>>>,
29}
30
31impl<FS> KubernetesProvider<FS>
32where
33    FS: FileSystem + Send + Sync + Clone,
34{
35    pub async fn new(filesystem: FS) -> Arc<Self> {
36        let k8s_client = KubernetesClient::new().await.unwrap();
37
38        Arc::new_cyclic(|weak| KubernetesProvider {
39            weak: weak.clone(),
40            capabilities: ProviderCapabilities {
41                requires_image: true,
42                has_resources: true,
43                prefix_with_full_path: false,
44                use_default_ports_in_cmd: true,
45            },
46            tmp_dir: std::env::temp_dir(),
47            k8s_client,
48            filesystem,
49            namespaces: RwLock::new(HashMap::new()),
50        })
51    }
52
53    pub fn tmp_dir(mut self, tmp_dir: impl Into<PathBuf>) -> Self {
54        self.tmp_dir = tmp_dir.into();
55        self
56    }
57}
58
59#[async_trait]
60impl<FS> Provider for KubernetesProvider<FS>
61where
62    FS: FileSystem + Send + Sync + Clone + 'static,
63{
64    fn name(&self) -> &str {
65        PROVIDER_NAME
66    }
67
68    fn capabilities(&self) -> &ProviderCapabilities {
69        &self.capabilities
70    }
71
72    async fn namespaces(&self) -> HashMap<String, DynNamespace> {
73        self.namespaces
74            .read()
75            .await
76            .iter()
77            .map(|(name, namespace)| (name.clone(), namespace.clone() as DynNamespace))
78            .collect()
79    }
80
81    async fn create_namespace(&self) -> Result<DynNamespace, ProviderError> {
82        let namespace = KubernetesNamespace::new(
83            &self.weak,
84            &self.tmp_dir,
85            &self.capabilities,
86            &self.k8s_client,
87            &self.filesystem,
88            None,
89        )
90        .await?;
91
92        self.namespaces
93            .write()
94            .await
95            .insert(namespace.name().to_string(), namespace.clone());
96
97        Ok(namespace)
98    }
99
100    async fn create_namespace_with_base_dir(
101        &self,
102        base_dir: &Path,
103    ) -> Result<DynNamespace, ProviderError> {
104        let namespace = KubernetesNamespace::new(
105            &self.weak,
106            &self.tmp_dir,
107            &self.capabilities,
108            &self.k8s_client,
109            &self.filesystem,
110            Some(base_dir),
111        )
112        .await?;
113
114        self.namespaces
115            .write()
116            .await
117            .insert(namespace.name().to_string(), namespace.clone());
118
119        Ok(namespace)
120    }
121
122    async fn create_namespace_from_json(
123        &self,
124        json_value: &serde_json::Value,
125    ) -> Result<DynNamespace, ProviderError> {
126        let (base_dir, name) = extract_namespace_info(json_value)?;
127
128        let namespace = KubernetesNamespace::attach_to_live(
129            &self.weak,
130            &self.capabilities,
131            &self.k8s_client,
132            &self.filesystem,
133            &base_dir,
134            &name,
135        )
136        .await?;
137
138        self.namespaces
139            .write()
140            .await
141            .insert(namespace.name().to_string(), namespace.clone());
142
143        Ok(namespace)
144    }
145}