From f5f32bfd8f705ac1ceec82455c6957e61b6ec484 Mon Sep 17 00:00:00 2001
From: Amanuel Engeda <74629455+engedaam@users.noreply.github.com>
Date: Mon, 27 Nov 2023 16:52:22 -0800
Subject: [PATCH] test: Fix Prefix Delegation enabling race condition (#5172)

---
 test/suites/scale/provisioning_test.go | 3 +++
 test/suites/scale/suite_test.go        | 4 ++--
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/test/suites/scale/provisioning_test.go b/test/suites/scale/provisioning_test.go
index 4f0d16a63dc6..c166031216d8 100644
--- a/test/suites/scale/provisioning_test.go
+++ b/test/suites/scale/provisioning_test.go
@@ -72,6 +72,9 @@ var _ = Describe("Provisioning", Label(debug.NoWatch), Label(debug.NoEvents), fu
 		// that will be allocated across this large number of nodes, despite the fact that the ENI CIDR space will
 		// be extremely under-utilized
 		env.ExpectPrefixDelegationDisabled()
+		DeferCleanup(func() {
+			env.ExpectPrefixDelegationEnabled()
+		})
 
 		replicasPerNode := 1
 		expectedNodeCount := 500
diff --git a/test/suites/scale/suite_test.go b/test/suites/scale/suite_test.go
index 8050d43cc98f..45d349cae332 100644
--- a/test/suites/scale/suite_test.go
+++ b/test/suites/scale/suite_test.go
@@ -30,20 +30,20 @@ func TestScale(t *testing.T) {
 	RegisterFailHandler(Fail)
 	BeforeSuite(func() {
 		env = aws.NewEnvironment(t)
+		env.ExpectPrefixDelegationEnabled()
 		SetDefaultEventuallyTimeout(time.Hour)
 	})
 	AfterSuite(func() {
+		env.ExpectPrefixDelegationDisabled()
 		env.Stop()
 	})
 	RunSpecs(t, "Scale")
 }
 
 var _ = BeforeEach(func() {
-	env.ExpectPrefixDelegationEnabled()
 	env.BeforeEach()
 })
 var _ = AfterEach(func() { env.Cleanup() })
 var _ = AfterEach(func() {
 	env.AfterEach()
-	env.ExpectPrefixDelegationDisabled()
 })