From 92740a25d4e2e97019a91fcd8c286b366247f83a Mon Sep 17 00:00:00 2001
From: Sean Malloy <sean.malloy@kohls.com>
Date: Tue, 23 Mar 2021 00:03:36 -0500
Subject: [PATCH] Add Initial Unit Tests For StatefulSets

---
 pkg/descheduler/evictions/evictions_test.go | 15 +++++++++++++++
 test/test_utils.go                          | 12 ++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/pkg/descheduler/evictions/evictions_test.go b/pkg/descheduler/evictions/evictions_test.go
index 20d770ce8..986a3247a 100644
--- a/pkg/descheduler/evictions/evictions_test.go
+++ b/pkg/descheduler/evictions/evictions_test.go
@@ -236,6 +236,21 @@ func TestIsEvictable(t *testing.T) {
 			evictLocalStoragePods: false,
 			priorityThreshold:     &lowPriority,
 			result:                true,
+		}, {
+			pod: test.BuildTestPod("p16", 400, 0, n1.Name, nil),
+			runBefore: func(pod *v1.Pod) {
+				pod.ObjectMeta.OwnerReferences = test.GetStatefulSetOwnerRefList()
+			},
+			evictLocalStoragePods: false,
+			result:                true,
+		}, {
+			pod: test.BuildTestPod("p17", 400, 0, n1.Name, nil),
+			runBefore: func(pod *v1.Pod) {
+				pod.Annotations = map[string]string{"descheduler.alpha.kubernetes.io/evict": "true"}
+				pod.ObjectMeta.OwnerReferences = test.GetStatefulSetOwnerRefList()
+			},
+			evictLocalStoragePods: false,
+			result:                true,
 		},
 	}
 
diff --git a/test/test_utils.go b/test/test_utils.go
index 206fa85c4..7b35ab8e0 100644
--- a/test/test_utils.go
+++ b/test/test_utils.go
@@ -79,6 +79,13 @@ func GetReplicaSetOwnerRefList() []metav1.OwnerReference {
 	return ownerRefList
 }
 
+// GetStatefulSetOwnerRefList returns the ownerRef needed for statefulset pod.
+func GetStatefulSetOwnerRefList() []metav1.OwnerReference {
+	ownerRefList := make([]metav1.OwnerReference, 0)
+	ownerRefList = append(ownerRefList, metav1.OwnerReference{Kind: "StatefulSet", APIVersion: "v1", Name: "statefulset-1"})
+	return ownerRefList
+}
+
 // GetDaemonSetOwnerRefList returns the ownerRef needed for daemonset pod.
 func GetDaemonSetOwnerRefList() []metav1.OwnerReference {
 	ownerRefList := make([]metav1.OwnerReference, 0)
@@ -142,6 +149,11 @@ func SetRSOwnerRef(pod *v1.Pod) {
 	pod.ObjectMeta.OwnerReferences = GetReplicaSetOwnerRefList()
 }
 
+// SetSSOwnerRef sets the given pod's owner to StatefulSet
+func SetSSOwnerRef(pod *v1.Pod) {
+	pod.ObjectMeta.OwnerReferences = GetStatefulSetOwnerRefList()
+}
+
 // SetDSOwnerRef sets the given pod's owner to DaemonSet
 func SetDSOwnerRef(pod *v1.Pod) {
 	pod.ObjectMeta.OwnerReferences = GetDaemonSetOwnerRefList()