Skip to main content

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    // Pre-existing namespace to spawn into, instead of creating one.
29    existing_namespace: Option<String>,
30    pub(super) namespaces: RwLock<HashMap<String, Arc<KubernetesNamespace<FS>>>>,
31}
32
33impl<FS> KubernetesProvider<FS>
34where
35    FS: FileSystem + Send + Sync + Clone,
36{
37    pub async fn new(filesystem: FS) -> Arc<Self> {
38        Self::build(filesystem, None).await
39    }
40
41    /// Spawns into `namespace`, which must already exist. Zombienet won't create it
42    /// nor delete it on drop, but `destroy()` still deletes it.
43    pub async fn new_in_namespace(filesystem: FS, namespace: impl Into<String>) -> Arc<Self> {
44        Self::build(filesystem, Some(namespace.into())).await
45    }
46
47    async fn build(filesystem: FS, existing_namespace: Option<String>) -> Arc<Self> {
48        let k8s_client = KubernetesClient::new().await.unwrap();
49        Self::with_client(filesystem, k8s_client, existing_namespace)
50    }
51
52    fn with_client(
53        filesystem: FS,
54        k8s_client: KubernetesClient,
55        existing_namespace: Option<String>,
56    ) -> Arc<Self> {
57        Arc::new_cyclic(|weak| KubernetesProvider {
58            weak: weak.clone(),
59            capabilities: ProviderCapabilities {
60                requires_image: true,
61                has_resources: true,
62                prefix_with_full_path: false,
63                use_default_ports_in_cmd: true,
64            },
65            tmp_dir: std::env::temp_dir(),
66            k8s_client,
67            filesystem,
68            existing_namespace,
69            namespaces: RwLock::new(HashMap::new()),
70        })
71    }
72
73    pub fn tmp_dir(mut self, tmp_dir: impl Into<PathBuf>) -> Self {
74        self.tmp_dir = tmp_dir.into();
75        self
76    }
77}
78
79#[async_trait]
80impl<FS> Provider for KubernetesProvider<FS>
81where
82    FS: FileSystem + Send + Sync + Clone + 'static,
83{
84    fn name(&self) -> &str {
85        PROVIDER_NAME
86    }
87
88    fn capabilities(&self) -> &ProviderCapabilities {
89        &self.capabilities
90    }
91
92    async fn namespaces(&self) -> HashMap<String, DynNamespace> {
93        self.namespaces
94            .read()
95            .await
96            .iter()
97            .map(|(name, namespace)| (name.clone(), namespace.clone() as DynNamespace))
98            .collect()
99    }
100
101    async fn create_namespace(&self) -> Result<DynNamespace, ProviderError> {
102        let namespace = KubernetesNamespace::new(
103            &self.weak,
104            &self.tmp_dir,
105            &self.capabilities,
106            &self.k8s_client,
107            &self.filesystem,
108            None,
109            self.existing_namespace.as_deref(),
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
121    async fn create_namespace_with_base_dir(
122        &self,
123        base_dir: &Path,
124    ) -> Result<DynNamespace, ProviderError> {
125        let namespace = KubernetesNamespace::new(
126            &self.weak,
127            &self.tmp_dir,
128            &self.capabilities,
129            &self.k8s_client,
130            &self.filesystem,
131            Some(base_dir),
132            self.existing_namespace.as_deref(),
133        )
134        .await?;
135
136        self.namespaces
137            .write()
138            .await
139            .insert(namespace.name().to_string(), namespace.clone());
140
141        Ok(namespace)
142    }
143
144    async fn create_namespace_from_json(
145        &self,
146        json_value: &serde_json::Value,
147    ) -> Result<DynNamespace, ProviderError> {
148        let (base_dir, name) = extract_namespace_info(json_value)?;
149
150        let namespace = KubernetesNamespace::attach_to_live(
151            &self.weak,
152            &self.capabilities,
153            &self.k8s_client,
154            &self.filesystem,
155            &base_dir,
156            &name,
157        )
158        .await?;
159
160        self.namespaces
161            .write()
162            .await
163            .insert(namespace.name().to_string(), namespace.clone());
164
165        Ok(namespace)
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use std::{collections::HashMap, ffi::OsString};
172
173    use mockito::{Matcher, Mock, Server, ServerGuard};
174    use support::fs::in_memory::{InMemoryFile, InMemoryFileSystem};
175
176    use super::*;
177
178    /// Provider talking to `server`. Unmocked requests get mockito's 501, so
179    /// spawning fails right after the namespace step. kube appends a query string,
180    /// hence `match_query` on exact-path mocks.
181    fn provider(
182        server: &ServerGuard,
183        existing_namespace: Option<&str>,
184    ) -> Arc<KubernetesProvider<InMemoryFileSystem>> {
185        // In-memory `create_dir` needs every ancestor of the tmp dir to exist. Going
186        // through `components()` drops the trailing `/` `$TMPDIR` may have.
187        let tmp_dir: PathBuf = std::env::temp_dir().components().collect();
188        let files: HashMap<OsString, InMemoryFile> = tmp_dir
189            .ancestors()
190            .map(|dir| (dir.as_os_str().to_owned(), InMemoryFile::dir()))
191            .collect();
192        let config = kube::Config::new(server.url().parse().unwrap());
193        let client = KubernetesClient::from_inner(kube::Client::try_from(config).unwrap());
194
195        KubernetesProvider::with_client(
196            InMemoryFileSystem::new(files),
197            client,
198            existing_namespace.map(String::from),
199        )
200    }
201
202    async fn mock_create_namespace(server: &mut ServerGuard, times: usize) -> Mock {
203        server
204            .mock("POST", "/api/v1/namespaces")
205            .match_query(Matcher::Any)
206            .with_status(201)
207            .with_body(r#"{"apiVersion":"v1","kind":"Namespace","metadata":{"name":"mock"}}"#)
208            .expect(times)
209            .create_async()
210            .await
211    }
212
213    async fn mock_delete_namespace(server: &mut ServerGuard, path: Matcher, times: usize) -> Mock {
214        server
215            .mock("DELETE", path)
216            .with_status(200)
217            .expect(times)
218            .create_async()
219            .await
220    }
221
222    #[tokio::test(flavor = "multi_thread")]
223    async fn existing_namespace_is_neither_created_nor_deleted() {
224        let mut server = Server::new_async().await;
225        let create = mock_create_namespace(&mut server, 0).await;
226        let delete = mock_delete_namespace(&mut server, Matcher::Any, 0).await;
227        let file_server = server
228            .mock("POST", "/api/v1/namespaces/adopted/pods")
229            .match_query(Matcher::Any)
230            .with_status(500)
231            .create_async()
232            .await;
233
234        let provider = provider(&server, Some("adopted"));
235        assert!(provider.create_namespace().await.is_err());
236        drop(provider);
237
238        file_server.assert_async().await;
239        create.assert_async().await;
240        delete.assert_async().await;
241    }
242
243    #[tokio::test(flavor = "multi_thread")]
244    async fn own_namespace_is_created_and_deleted_on_drop() {
245        let mut server = Server::new_async().await;
246        let create = mock_create_namespace(&mut server, 1).await;
247        let delete = mock_delete_namespace(
248            &mut server,
249            Matcher::Regex("^/api/v1/namespaces/zombie-".into()),
250            1,
251        )
252        .await;
253
254        let provider = provider(&server, None);
255        assert!(provider.create_namespace().await.is_err());
256        drop(provider);
257
258        create.assert_async().await;
259        delete.assert_async().await;
260    }
261}