From d45790153257814f06cf2c722ef63022f323fdba Mon Sep 17 00:00:00 2001
From: Reed Schalo <reed0schalo@gmail.com>
Date: Thu, 26 Sep 2024 16:34:33 -0700
Subject: [PATCH] fix: update test to reduce create fleet payload size (#7087)

---
 test/suites/integration/termination_test.go   |  2 +-
 .../nodeclaim/garbage_collection_test.go      |  2 +-
 test/suites/nodeclaim/nodeclaim_test.go       | 25 +++++++++++++++++--
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/test/suites/integration/termination_test.go b/test/suites/integration/termination_test.go
index cbd32fb51a27..a42de19966aa 100644
--- a/test/suites/integration/termination_test.go
+++ b/test/suites/integration/termination_test.go
@@ -44,7 +44,7 @@ var _ = Describe("Termination", func() {
 		env.ExpectDeleted(nodes[0])
 		env.EventuallyExpectNotFound(nodes[0])
 		Eventually(func(g Gomega) {
-			g.Expect(lo.FromPtr(env.GetInstanceByID(instanceID).State.Name)).To(Equal("shutting-down"))
+			g.Expect(lo.FromPtr(env.GetInstanceByID(instanceID).State.Name)).To(BeElementOf("shutting-down", "terminated"))
 		}, time.Second*10).Should(Succeed())
 	})
 })
diff --git a/test/suites/nodeclaim/garbage_collection_test.go b/test/suites/nodeclaim/garbage_collection_test.go
index 35d5dd4700bc..186fe8cfd90d 100644
--- a/test/suites/nodeclaim/garbage_collection_test.go
+++ b/test/suites/nodeclaim/garbage_collection_test.go
@@ -140,7 +140,7 @@ var _ = Describe("GarbageCollection", func() {
 		// Eventually expect the node and the instance to be removed (shutting-down)
 		env.EventuallyExpectNotFound(node)
 		Eventually(func(g Gomega) {
-			g.Expect(lo.FromPtr(env.GetInstanceByID(aws.StringValue(out.Instances[0].InstanceId)).State.Name)).To(Equal("shutting-down"))
+			g.Expect(lo.FromPtr(env.GetInstanceByID(aws.StringValue(out.Instances[0].InstanceId)).State.Name)).To(BeElementOf("shutting-down", "terminated"))
 		}, time.Second*10).Should(Succeed())
 	})
 	It("should succeed to garbage collect an Instance that was deleted without the cluster's knowledge", func() {
diff --git a/test/suites/nodeclaim/nodeclaim_test.go b/test/suites/nodeclaim/nodeclaim_test.go
index 82ec0448d0dd..a561c798ca72 100644
--- a/test/suites/nodeclaim/nodeclaim_test.go
+++ b/test/suites/nodeclaim/nodeclaim_test.go
@@ -48,6 +48,13 @@ var _ = Describe("StandaloneNodeClaim", func() {
 							Values:   []string{"c"},
 						},
 					},
+					{
+						NodeSelectorRequirement: v1.NodeSelectorRequirement{
+							Key:      v1.LabelArchStable,
+							Operator: v1.NodeSelectorOpIn,
+							Values:   []string{"arm64"},
+						},
+					},
 					{
 						NodeSelectorRequirement: v1.NodeSelectorRequirement{
 							Key:      corev1beta1.CapacityTypeLabelKey,
@@ -201,6 +208,13 @@ var _ = Describe("StandaloneNodeClaim", func() {
 							Values:   []string{"c"},
 						},
 					},
+					{
+						NodeSelectorRequirement: v1.NodeSelectorRequirement{
+							Key:      v1.LabelArchStable,
+							Operator: v1.NodeSelectorOpIn,
+							Values:   []string{"arm64"},
+						},
+					},
 					{
 						NodeSelectorRequirement: v1.NodeSelectorRequirement{
 							Key:      corev1beta1.CapacityTypeLabelKey,
@@ -226,7 +240,7 @@ var _ = Describe("StandaloneNodeClaim", func() {
 		env.EventuallyExpectNotFound(nodeClaim, node)
 
 		Eventually(func(g Gomega) {
-			g.Expect(lo.FromPtr(env.GetInstanceByID(instanceID).State.Name)).To(Equal("shutting-down"))
+			g.Expect(lo.FromPtr(env.GetInstanceByID(instanceID).State.Name)).To(BeElementOf("shutting-down", "terminated"))
 		}, time.Second*10).Should(Succeed())
 	})
 	It("should delete a NodeClaim from the node termination finalizer", func() {
@@ -240,6 +254,13 @@ var _ = Describe("StandaloneNodeClaim", func() {
 							Values:   []string{"c"},
 						},
 					},
+					{
+						NodeSelectorRequirement: v1.NodeSelectorRequirement{
+							Key:      v1.LabelArchStable,
+							Operator: v1.NodeSelectorOpIn,
+							Values:   []string{"arm64"},
+						},
+					},
 					{
 						NodeSelectorRequirement: v1.NodeSelectorRequirement{
 							Key:      corev1beta1.CapacityTypeLabelKey,
@@ -265,7 +286,7 @@ var _ = Describe("StandaloneNodeClaim", func() {
 		env.EventuallyExpectNotFound(nodeClaim, node)
 
 		Eventually(func(g Gomega) {
-			g.Expect(lo.FromPtr(env.GetInstanceByID(instanceID).State.Name)).To(Equal("shutting-down"))
+			g.Expect(lo.FromPtr(env.GetInstanceByID(instanceID).State.Name)).To(BeElementOf("shutting-down", "terminated"))
 		}, time.Second*10).Should(Succeed())
 	})
 	It("should create a NodeClaim with custom labels passed through the userData", func() {