From 354f692cd842fac59c791907753a530ef1aad9f1 Mon Sep 17 00:00:00 2001
From: Nick Tran <10810510+njtran@users.noreply.github.com>
Date: Mon, 5 Feb 2024 14:59:59 -0800
Subject: [PATCH] test: fix drift/expiration test flakes (#5601)

---
 test/suites/drift/suite_test.go      |  9 ++-------
 test/suites/expiration/suite_test.go | 13 ++++---------
 2 files changed, 6 insertions(+), 16 deletions(-)

diff --git a/test/suites/drift/suite_test.go b/test/suites/drift/suite_test.go
index 032880c65680..fe21255c0260 100644
--- a/test/suites/drift/suite_test.go
+++ b/test/suites/drift/suite_test.go
@@ -218,6 +218,8 @@ var _ = Describe("Drift", func() {
 			dep.Spec.Replicas = lo.ToPtr[int32](3)
 			env.ExpectUpdated(dep)
 
+			// First expect there to be 3 pods, then try to spread the pods.
+			env.EventuallyExpectHealthyPodCount(selector, 3)
 			env.ForcePodsToSpread(nodes...)
 			env.EventuallyExpectHealthyPodCount(selector, 3)
 
@@ -226,8 +228,6 @@ var _ = Describe("Drift", func() {
 			for _, node := range nodes {
 				Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(node), node)).To(Succeed())
 				node.Finalizers = append(node.Finalizers, common.TestingFinalizer)
-				// Set nodes as unschedulable so that pod nomination doesn't delay disruption for the second disruption action
-				node.Spec.Unschedulable = true
 				env.ExpectUpdated(node)
 			}
 
@@ -246,11 +246,6 @@ var _ = Describe("Drift", func() {
 				env.ExpectUpdated(pod)
 			}
 
-			// Mark one node as schedulable so the other two nodes can schedule to this node and delete.
-			Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(nodes[0]), nodes[0])).To(Succeed())
-			nodes[0].Spec.Unschedulable = false
-			env.ExpectUpdated(nodes[0])
-
 			// Ensure that we get two nodes tainted, and they have overlap during the drift
 			env.EventuallyExpectTaintedNodeCount("==", 2)
 			nodes = env.ConsistentlyExpectDisruptionsWithNodeCount(2, 3, 30*time.Second)
diff --git a/test/suites/expiration/suite_test.go b/test/suites/expiration/suite_test.go
index accf8200d2db..4ba4290a14b8 100644
--- a/test/suites/expiration/suite_test.go
+++ b/test/suites/expiration/suite_test.go
@@ -297,6 +297,8 @@ var _ = Describe("Expiration", func() {
 			dep.Spec.Replicas = lo.ToPtr[int32](3)
 			env.ExpectUpdated(dep)
 
+			// First expect there to be 3 pods, then try to spread the pods.
+			env.EventuallyExpectHealthyPodCount(selector, 3)
 			env.ForcePodsToSpread(nodes...)
 			env.EventuallyExpectHealthyPodCount(selector, 3)
 
@@ -305,8 +307,6 @@ var _ = Describe("Expiration", func() {
 			for _, node := range nodes {
 				Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(node), node)).To(Succeed())
 				node.Finalizers = append(node.Finalizers, common.TestingFinalizer)
-				// Set nodes as unschedulable so that pod nomination doesn't delay disruption for the second disruption action
-				node.Spec.Unschedulable = true
 				env.ExpectUpdated(node)
 			}
 
@@ -325,11 +325,6 @@ var _ = Describe("Expiration", func() {
 				env.ExpectUpdated(pod)
 			}
 
-			// Mark one node as schedulable so the other two nodes can schedule to this node and delete.
-			Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(nodes[0]), nodes[0])).To(Succeed())
-			nodes[0].Spec.Unschedulable = false
-			env.ExpectUpdated(nodes[0])
-
 			// Ensure that we get two nodes tainted, and they have overlap during the expiration
 			env.EventuallyExpectTaintedNodeCount("==", 2)
 			nodes = env.ConsistentlyExpectDisruptionsWithNodeCount(2, 3, 5*time.Second)
@@ -475,7 +470,7 @@ var _ = Describe("Expiration", func() {
 			env.ExpectUpdated(pod)
 		}
 
-		// Eventually the node will be set as unschedulable, which means its actively being deprovisioned
+		// Eventually the node will be tainted, which means its actively being disrupted
 		Eventually(func(g Gomega) {
 			g.Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(node), node)).Should(Succeed())
 			_, ok := lo.Find(node.Spec.Taints, func(t v1.Taint) bool {
@@ -539,7 +534,7 @@ var _ = Describe("Expiration", func() {
 			env.ExpectUpdated(pod)
 		}
 
-		// Eventually the node will be set as unschedulable, which means its actively being deprovisioned
+		// Eventually the node will be tainted, which means its actively being disrupted
 		Eventually(func(g Gomega) {
 			g.Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(node), node)).Should(Succeed())
 			_, ok := lo.Find(node.Spec.Taints, func(t v1.Taint) bool {