From e9628355d975a8bb4fb50a4ec8efa399f3fd2f44 Mon Sep 17 00:00:00 2001
From: rrajesh <rrajesh@cloudera.com>
Date: Tue, 5 Nov 2024 17:09:39 +0530
Subject: [PATCH] [YUNIKORN-2963] E2E Test for Foreign Pod Tracking

---
 test/e2e/foreign_pod/foreign_pod_test.go | 113 +++++++++++++++++++++++
 1 file changed, 113 insertions(+)

diff --git a/test/e2e/foreign_pod/foreign_pod_test.go b/test/e2e/foreign_pod/foreign_pod_test.go
index 932afbfd5..c95e06a78 100644
--- a/test/e2e/foreign_pod/foreign_pod_test.go
+++ b/test/e2e/foreign_pod/foreign_pod_test.go
@@ -71,6 +71,119 @@ var _ = Describe("", func() {
 			Ω(foreignNodes[uid]).To(Equal(kubeNodes[uid]), "pod %s is tracked under incorrect node", uid)
 		}
 	})
+	It("Verify foreign pod tracking fails for untracked or incorrectly mapped pods", func() {
+		By("Retrieving foreign pods from kube-system")
+		kClient = k8s.KubeCtl{}
+		Ω(kClient.SetClient()).To(BeNil())
+		podList, err := kClient.GetPods(kubeSystem)
+		Ω(err).NotTo(gomega.HaveOccurred())
+
+		kubeUIDs := make(map[string]bool)
+		kubeNodes := make(map[string]string)
+		for _, pod := range podList.Items {
+			kubeUIDs[string(pod.UID)] = true
+			kubeNodes[string(pod.UID)] = pod.Spec.NodeName
+			fmt.Fprintf(ginkgo.GinkgoWriter, "pod: %s, uid: %s, node: %s\n", pod.Name, pod.UID, pod.Spec.NodeName)
+		}
+
+		// Simulate an error in foreign pod tracking
+		By("Simulating incorrect foreign allocation tracking")
+		var restClient yunikorn.RClient
+		nodes, err := restClient.GetNodes("default")
+		Ω(err).NotTo(gomega.HaveOccurred())
+
+		foreignAllocs := make(map[string]bool)
+		foreignNodes := make(map[string]string)
+		for _, n := range *nodes {
+			fmt.Fprintf(ginkgo.GinkgoWriter, "Checking node %s\n", n.NodeID)
+			if len(n.ForeignAllocations) > 0 {
+				for _, falloc := range n.ForeignAllocations {
+					// Introduce a mismatch to simulate an error
+					if string(falloc.AllocationKey) == "incorrect-uid" {
+						foreignAllocs[falloc.AllocationKey] = true
+						foreignNodes[falloc.AllocationKey] = "wrong-node"
+					} else {
+						foreignAllocs[falloc.AllocationKey] = true
+						foreignNodes[falloc.AllocationKey] = falloc.NodeID
+					}
+				}
+			}
+		}
+
+		// Check that an incorrect UID or mapping triggers an error
+		for uid := range kubeUIDs {
+			if uid == "incorrect-uid" {
+				Ω(foreignAllocs[uid]).To(Equal(false), "Simulated error: pod %s from kube-system should not be tracked", uid)
+				Ω(foreignNodes[uid]).To(Equal("wrong-node"), "Simulated error: pod %s should be tracked under incorrect node", uid)
+			} else {
+				Ω(foreignAllocs[uid]).To(Equal(true), "pod %s from kube-system is not tracked in Yunikorn", uid)
+				Ω(foreignNodes[uid]).To(Equal(kubeNodes[uid]), "pod %s is tracked under incorrect node", uid)
+			}
+		}
+	})
+	It("Verify multiple foreign pods", func() {
+		By("Retrieving foreign pods from kube-system")
+		kClient = k8s.KubeCtl{}
+		Ω(kClient.SetClient()).To(BeNil())
+		podList, err := kClient.GetPods(kubeSystem)
+		Ω(err).NotTo(gomega.HaveOccurred())
+
+		kubeUIDs := make(map[string]bool)
+		for _, pod := range podList.Items {
+			kubeUIDs[string(pod.UID)] = true
+		}
+
+		// retrieve foreign pod info
+		By("Retrieving foreign allocations")
+		var restClient yunikorn.RClient
+		nodes, err := restClient.GetNodes("default")
+		Ω(err).NotTo(gomega.HaveOccurred())
+		foreignAllocs := make(map[string]bool)
+		for _, n := range *nodes {
+			for _, falloc := range n.ForeignAllocations {
+				foreignAllocs[falloc.AllocationKey] = true
+			}
+		}
+
+		// check that all UIDs from kube-system are tracked properly
+		for uid := range kubeUIDs {
+			Ω(foreignAllocs[uid]).To(Equal(true), "pod %s from kube-system is not tracked in Yunikorn", uid)
+		}
+	})
+	It("Verify foreign pods on different nodes", func() {
+		By("Retrieving foreign pods from kube-system")
+		kClient = k8s.KubeCtl{}
+		Ω(kClient.SetClient()).To(BeNil())
+		podList, err := kClient.GetPods(kubeSystem)
+		Ω(err).NotTo(gomega.HaveOccurred())
+
+		kubeUIDs := make(map[string]bool)
+		kubeNodes := make(map[string]string)
+		for _, pod := range podList.Items {
+			kubeUIDs[string(pod.UID)] = true
+			kubeNodes[string(pod.UID)] = pod.Spec.NodeName
+		}
+
+		// retrieve foreign pod info
+		By("Retrieving foreign allocations")
+		var restClient yunikorn.RClient
+		nodes, err := restClient.GetNodes("default")
+		Ω(err).NotTo(gomega.HaveOccurred())
+		foreignAllocs := make(map[string]bool)
+		foreignNodes := make(map[string]string)
+		for _, n := range *nodes {
+			for _, falloc := range n.ForeignAllocations {
+				foreignAllocs[falloc.AllocationKey] = true
+				foreignNodes[falloc.AllocationKey] = falloc.NodeID
+			}
+		}
+
+		// check that all UIDs from kube-system are tracked properly
+		for uid := range kubeUIDs {
+			Ω(foreignAllocs[uid]).To(Equal(true), "pod %s from kube-system is not tracked in Yunikorn", uid)
+			Ω(foreignNodes[uid]).To(Equal(kubeNodes[uid]), "pod %s is tracked under incorrect node", uid)
+		}
+	})
 
 	ginkgo.AfterEach(func() {
 		tests.DumpClusterInfoIfSpecFailed(suiteName, []string{kubeSystem})