Unit test

linkerd · Dec 28, 2023 · f3c5459 · f3c5459
1 parent cc72fe8
commit f3c5459
Show file tree

Hide file tree

Showing 3 changed files with 90 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/cni-repair-controller/Cargo.toml b/cni-repair-controller/Cargo.toml
@@ -15,7 +15,7 @@ rustls-tls = ["kube/rustls-tls"]
 
 [dependencies]
 anyhow = "1"
-futures-util = "0.3"
+futures = "0.3"
 k8s-openapi = { version = "0.20", features = ["v1_22"] }
 kube = { version = "0.87", features = ["runtime"] }
 openssl = { version = "0.10.60", optional = true }
@@ -35,3 +35,6 @@ features = ["admin", "clap", "prometheus-client", "runtime"]
 [dependencies.tokio]
 version = "1"
 features = ["macros", "parking_lot", "rt", "rt-multi-thread"]
+
+[dev-dependencies]
+chrono = "0.4"
diff --git a/cni-repair-controller/src/lib.rs b/cni-repair-controller/src/lib.rs
@@ -1,4 +1,4 @@
-use futures_util::{Stream, StreamExt};
+use futures::{Stream, StreamExt};
 use k8s_openapi::api::core::v1::{ObjectReference, Pod};
 use kube::{
     runtime::{
@@ -163,3 +163,84 @@ impl Metrics {
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use chrono::Utc;
+    use k8s_openapi::api::core::v1::{
+        ContainerState, ContainerStateTerminated, ContainerStatus, PodStatus,
+    };
+    use k8s_openapi::apimachinery::pkg::apis::meta::v1::{ObjectMeta, Time};
+    use tokio::{
+        sync::mpsc::error::TryRecvError,
+        time::{self, Duration},
+    };
+
+    #[tokio::test]
+    async fn test_process_events() {
+        let mut prom = prometheus_client::registry::Registry::default();
+        let metrics =
+            Metrics::register(prom.sub_registry_with_prefix("linkerd_cni_repair_controller"));
+
+        // This pod should be ignored
+        let pod1 = Pod {
+            metadata: ObjectMeta {
+                name: Some("pod1".to_string()),
+                namespace: Some("default".to_string()),
+                ..Default::default()
+            },
+            ..Default::default()
+        };
+
+        // This pod should be processed
+        let pod2 = Pod {
+            metadata: ObjectMeta {
+                name: Some("pod2".to_string()),
+                namespace: Some("default".to_string()),
+                ..Default::default()
+            },
+            status: Some(PodStatus {
+                init_container_statuses: Some(vec![ContainerStatus {
+                    name: "linkerd-network-validator".to_string(),
+                    last_state: Some(ContainerState {
+                        terminated: Some(ContainerStateTerminated {
+                            exit_code: UNSUCCESSFUL_EXIT_CODE,
+                            ..Default::default()
+                        }),
+                        ..Default::default()
+                    }),
+                    ..Default::default()
+                }]),
+                ..Default::default()
+            }),
+            ..Default::default()
+        };
+
+        // This pod should be ignored
+        let pod3 = Pod {
+            metadata: ObjectMeta {
+                name: Some("pod2".to_string()),
+                namespace: Some("default".to_string()),
+                deletion_timestamp: Some(Time(Utc::now())),
+                ..Default::default()
+            },
+            ..pod2.clone()
+        };
+
+        let (tx, mut rx) = mpsc::channel(EVENT_CHANNEL_CAPACITY);
+        let stream = futures::stream::iter(vec![
+            watcher::Event::Applied(pod1),
+            watcher::Event::Applied(pod2),
+            watcher::Event::Applied(pod3),
+        ]);
+
+        tokio::spawn(process_events(stream, tx, metrics));
+        time::sleep(Duration::from_secs(2)).await;
+        let msg = rx.try_recv();
+        let object_ref = msg.unwrap();
+        assert_eq!(object_ref.name, Some("pod2".to_string()));
+        let msg = rx.try_recv();
+        assert_eq!(msg, Err(TryRecvError::Disconnected));
+    }
+}