Add fix for nodes leaking into subsequent tests for budgets

test/pkg/environment/common/expectations.go

@@ -739,7 +739,7 @@ func (env *Environment) printControllerLogs(options *v1.PodLogOptions) {
 		temp := options.DeepCopy() // local version of the log options
 
 		fmt.Printf("------- pod/%s -------\n", pod.Name)
-		if pod.Status.ContainerStatuses[0].RestartCount > 0 {
+		if len(pod.Status.ContainerStatuses) > 0 && pod.Status.ContainerStatuses[0].RestartCount > 0 {
 			fmt.Printf("[PREVIOUS CONTAINER LOGS]\n")
 			temp.Previous = true
 		}

test/suites/consolidation/suite_test.go

@@ -283,19 +283,12 @@ var _ = Describe("Consolidation", func() {
 
 			env.ExpectCreated(nodeClass, nodePool, deployments[0], deployments[1], deployments[2], deployments[3], deployments[4])
 
-			env.EventuallyExpectCreatedNodeClaimCount("==", 5)
-			nodes := env.EventuallyExpectCreatedNodeCount("==", 5)
+			originalNodeClaims := env.EventuallyExpectCreatedNodeClaimCount("==", 5)
+			originalNodes := env.EventuallyExpectCreatedNodeCount("==", 5)
+
 			// Check that all daemonsets and deployment pods are online
 			env.EventuallyExpectHealthyPodCount(selector, int(numPods)*2)
 
-			By("cordoning and adding finalizer to the nodes")
-			// Add a finalizer to each node so that we can stop termination disruptions
-			for _, node := range nodes {
-				Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(node), node)).To(Succeed())
-				node.Finalizers = append(node.Finalizers, common.TestingFinalizer)
-				env.ExpectUpdated(node)
-			}
-
 			// Delete the daemonset so that the nodes can be consolidated to smaller size
 			env.ExpectDeleted(ds)
 			// Check that all daemonsets and deployment pods are online
@@ -309,12 +302,14 @@ var _ = Describe("Consolidation", func() {
 			env.EventuallyExpectTaintedNodeCount("==", 3)
 			env.EventuallyExpectNodeClaimCount("==", 8)
 			env.EventuallyExpectNodeCount("==", 8)
-			nodes = env.ConsistentlyExpectDisruptionsWithNodeCount(3, 8, 5*time.Second)
-
-			for _, node := range nodes {
-				Expect(env.ExpectTestingFinalizerRemoved(node)).To(Succeed())
-			}
-			env.EventuallyExpectNotFound(nodes[0], nodes[1], nodes[2])
+			env.ConsistentlyExpectDisruptionsWithNodeCount(3, 8, 5*time.Second)
+
+			// Eventually expect all the nodes to be rolled and completely removed
+			// Since this completes the disruption operation, this also ensures that we aren't leaking nodes into subsequent
+			// tests since nodeclaims that are actively replacing but haven't brought-up nodes yet can register nodes later
+			env.EventuallyExpectNotFound(lo.Map(originalNodes, func(n *v1.Node, _ int) client.Object { return n })...)
+			env.EventuallyExpectNotFound(lo.Map(originalNodeClaims, func(n *corev1beta1.NodeClaim, _ int) client.Object { return n })...)
+			env.ExpectNodeClaimCount("==", 5)
 			env.ExpectNodeCount("==", 5)
 		})
 		It("should not allow consolidation if the budget is fully blocking", func() {

test/suites/drift/suite_test.go

@@ -97,8 +97,6 @@ var _ = Describe("Drift", func() {
 			},
 		})
 		selector = labels.SelectorFromSet(dep.Spec.Selector.MatchLabels)
-
-		env.ExpectSettingsOverridden(v1.EnvVar{Name: "FEATURE_GATES", Value: "Drift=true"})
 	})
 	Context("Budgets", func() {
 		It("should respect budgets for empty drift", func() {
@@ -304,36 +302,28 @@ var _ = Describe("Drift", func() {
 
 			env.ExpectCreated(nodeClass, nodePool, deployments[0], deployments[1], deployments[2], deployments[3], deployments[4])
 
-			env.EventuallyExpectCreatedNodeClaimCount("==", 5)
-			nodes := env.EventuallyExpectCreatedNodeCount("==", 5)
-			// Check that all deployment pods are online
-			env.EventuallyExpectHealthyPodCount(selector, numPods)
+			originalNodeClaims := env.EventuallyExpectCreatedNodeClaimCount("==", 5)
+			originalNodes := env.EventuallyExpectCreatedNodeCount("==", 5)
 
-			By("cordoning and adding finalizer to the nodes")
-			// Add a finalizer to each node so that we can stop termination disruptions
-			for _, node := range nodes {
-				Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(node), node)).To(Succeed())
-				node.Finalizers = append(node.Finalizers, common.TestingFinalizer)
-				env.ExpectUpdated(node)
-			}
-
-			// Check that all daemonsets and deployment pods are online
+			// Check that all deployment pods are online
 			env.EventuallyExpectHealthyPodCount(selector, numPods)
 
 			By("drifting the nodepool")
 			nodePool.Spec.Template.Annotations = lo.Assign(nodePool.Spec.Template.Annotations, map[string]string{"test-annotation": "drift"})
 			env.ExpectUpdated(nodePool)
 
-			// Ensure that we get two nodes tainted, and they have overlap during the drift
+			// Ensure that we get three nodes tainted, and they have overlap during the drift
 			env.EventuallyExpectTaintedNodeCount("==", 3)
 			env.EventuallyExpectNodeClaimCount("==", 8)
 			env.EventuallyExpectNodeCount("==", 8)
-			nodes = env.ConsistentlyExpectDisruptionsWithNodeCount(3, 8, 5*time.Second)
-
-			for _, node := range nodes {
-				Expect(env.ExpectTestingFinalizerRemoved(node)).To(Succeed())
-			}
-			env.EventuallyExpectNotFound(nodes[0], nodes[1], nodes[2])
+			env.ConsistentlyExpectDisruptionsWithNodeCount(3, 8, 5*time.Second)
+
+			// Eventually expect all the nodes to be rolled and completely removed
+			// Since this completes the disruption operation, this also ensures that we aren't leaking nodes into subsequent
+			// tests since nodeclaims that are actively replacing but haven't brought-up nodes yet can register nodes later
+			env.EventuallyExpectNotFound(lo.Map(originalNodes, func(n *v1.Node, _ int) client.Object { return n })...)
+			env.EventuallyExpectNotFound(lo.Map(originalNodeClaims, func(n *corev1beta1.NodeClaim, _ int) client.Object { return n })...)
+			env.ExpectNodeClaimCount("==", 5)
 			env.ExpectNodeCount("==", 5)
 		})
 		It("should not allow drift if the budget is fully blocking", func() {

test/suites/expiration/suite_test.go

@@ -384,18 +384,11 @@ var _ = Describe("Expiration", func() {
 			env.ExpectCreated(nodeClass, nodePool, deployments[0], deployments[1], deployments[2], deployments[3], deployments[4])
 
 			env.EventuallyExpectCreatedNodeClaimCount("==", 5)
-			nodes := env.EventuallyExpectCreatedNodeCount("==", 5)
+			env.EventuallyExpectCreatedNodeCount("==", 5)
+
 			// Check that all daemonsets and deployment pods are online
 			env.EventuallyExpectHealthyPodCount(selector, numPods)
 
-			By("cordoning and adding finalizer to the nodes")
-			// Add a finalizer to each node so that we can stop termination disruptions
-			for _, node := range nodes {
-				Expect(env.Client.Get(env.Context, client.ObjectKeyFromObject(node), node)).To(Succeed())
-				node.Finalizers = append(node.Finalizers, common.TestingFinalizer)
-				env.ExpectUpdated(node)
-			}
-
 			// Check that all daemonsets and deployment pods are online
 			env.EventuallyExpectHealthyPodCount(selector, numPods)
 
@@ -407,19 +400,16 @@ var _ = Describe("Expiration", func() {
 			env.EventuallyExpectTaintedNodeCount("==", 3)
 			env.EventuallyExpectNodeClaimCount("==", 8)
 			env.EventuallyExpectNodeCount("==", 8)
-			nodes = env.ConsistentlyExpectDisruptionsWithNodeCount(3, 8, 5*time.Second)
+			nodes := env.ConsistentlyExpectDisruptionsWithNodeCount(3, 8, 5*time.Second)
 
 			// Set the expireAfter to "Never" to make sure new node isn't deleted
 			// This is CRITICAL since it prevents nodes that are immediately spun up from immediately being expired and
 			// racing at the end of the E2E test, leaking node resources into subsequent tests
 			nodePool.Spec.Disruption.ExpireAfter.Duration = nil
 			env.ExpectUpdated(nodePool)
 
-			for _, node := range nodes {
-				Expect(env.ExpectTestingFinalizerRemoved(node)).To(Succeed())
-			}
-
 			env.EventuallyExpectNotFound(nodes[0], nodes[1], nodes[2])
+			env.ExpectNodeClaimCount("==", 5)
 			env.ExpectNodeCount("==", 5)
 		})
 		It("should not allow expiration if the budget is fully blocking", func() {