Refactor featureGate usage in unit tests

kubernetes-sigs · Aug 7, 2024 · abe85cb · abe85cb
1 parent 444f896
commit abe85cb
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 39 deletions.
diff --git a/api/v1beta1/azuremanagedcontrolplane_webhook_test.go b/api/v1beta1/azuremanagedcontrolplane_webhook_test.go
@@ -1662,29 +1662,31 @@ func TestAzureManagedControlPlane_ValidateCreate(t *testing.T) {
 
 func TestAzureManagedControlPlane_ValidateCreateFailure(t *testing.T) {
 	tests := []struct {
-		name        string
-		amcp        *AzureManagedControlPlane
-		deferFunc   func()
-		expectError bool
+		name               string
+		amcp               *AzureManagedControlPlane
+		featureGateEnabled *bool
+		expectError        bool
 	}{
 		{
-			name:        "feature gate explicitly disabled",
-			amcp:        getKnownValidAzureManagedControlPlane(),
-			deferFunc:   utilfeature.SetFeatureGateDuringTest(t, feature.Gates, capifeature.MachinePool, false),
-			expectError: true,
+			name:               "feature gate explicitly disabled",
+			amcp:               getKnownValidAzureManagedControlPlane(),
+			featureGateEnabled: ptr.To(false),
+			expectError:        true,
 		},
 		{
-			name:        "feature gate implicitly enabled",
-			amcp:        getKnownValidAzureManagedControlPlane(),
-			deferFunc:   func() {},
-			expectError: false,
+			name:               "feature gate implicitly enabled",
+			amcp:               getKnownValidAzureManagedControlPlane(),
+			featureGateEnabled: nil,
+			expectError:        false,
 		},
 	}
 	client := mockClient{ReturnError: false}
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
 			g := NewWithT(t)
-			defer tc.deferFunc()
+			if tc.featureGateEnabled != nil {
+				defer utilfeature.SetFeatureGateDuringTest(t, feature.Gates, capifeature.MachinePool, *tc.featureGateEnabled)()
+			}
 			mcpw := &azureManagedControlPlaneWebhook{
 				Client: client,
 			}

diff --git a/api/v1beta1/azuremanagedmachinepool_webhook_test.go b/api/v1beta1/azuremanagedmachinepool_webhook_test.go
@@ -1307,27 +1307,29 @@ func TestAzureManagedMachinePool_ValidateCreate(t *testing.T) {
 
 func TestAzureManagedMachinePool_ValidateCreateFailure(t *testing.T) {
 	tests := []struct {
-		name        string
-		ammp        *AzureManagedMachinePool
-		deferFunc   func()
-		expectError bool
+		name               string
+		ammp               *AzureManagedMachinePool
+		featureGateEnabled *bool
+		expectError        bool
 	}{
 		{
-			name:        "feature gate explicitly disabled",
-			ammp:        getKnownValidAzureManagedMachinePool(),
-			deferFunc:   utilfeature.SetFeatureGateDuringTest(t, feature.Gates, capifeature.MachinePool, false),
-			expectError: true,
+			name:               "feature gate explicitly disabled",
+			ammp:               getKnownValidAzureManagedMachinePool(),
+			featureGateEnabled: ptr.To(false),
+			expectError:        true,
 		},
 		{
-			name:        "feature gate implicitly enabled",
-			ammp:        getKnownValidAzureManagedMachinePool(),
-			deferFunc:   func() {},
-			expectError: false,
+			name:               "feature gate implicitly enabled",
+			ammp:               getKnownValidAzureManagedMachinePool(),
+			featureGateEnabled: nil,
+			expectError:        false,
 		},
 	}
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			defer tc.deferFunc()
+			if tc.featureGateEnabled != nil {
+				defer utilfeature.SetFeatureGateDuringTest(t, feature.Gates, capifeature.MachinePool, *tc.featureGateEnabled)()
+			}
 			g := NewWithT(t)
 			mw := &azureManagedMachinePoolWebhook{}
 			_, err := mw.ValidateCreate(context.Background(), tc.ammp)

diff --git a/exp/api/v1beta1/azuremachinepool_webhook_test.go b/exp/api/v1beta1/azuremachinepool_webhook_test.go
@@ -700,27 +700,29 @@ func TestAzureMachinePool_ValidateCreateFailure(t *testing.T) {
 	g := NewWithT(t)
 
 	tests := []struct {
-		name        string
-		amp         *AzureMachinePool
-		deferFunc   func()
-		expectError bool
+		name               string
+		amp                *AzureMachinePool
+		featureGateEnabled *bool
+		expectError        bool
 	}{
 		{
-			name:        "feature gate explicitly disabled",
-			amp:         getKnownValidAzureMachinePool(),
-			deferFunc:   utilfeature.SetFeatureGateDuringTest(t, feature.Gates, capifeature.MachinePool, false),
-			expectError: true,
+			name:               "feature gate explicitly disabled",
+			amp:                getKnownValidAzureMachinePool(),
+			featureGateEnabled: ptr.To(false),
+			expectError:        true,
 		},
 		{
-			name:        "feature gate implicitly enabled",
-			amp:         getKnownValidAzureMachinePool(),
-			deferFunc:   func() {},
-			expectError: false,
+			name:               "feature gate implicitly enabled",
+			amp:                getKnownValidAzureMachinePool(),
+			featureGateEnabled: nil,
+			expectError:        false,
 		},
 	}
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			defer tc.deferFunc()
+			if tc.featureGateEnabled != nil {
+				defer utilfeature.SetFeatureGateDuringTest(t, feature.Gates, capifeature.MachinePool, *tc.featureGateEnabled)()
+			}
 			ampw := &azureMachinePoolWebhook{}
 			_, err := ampw.ValidateCreate(context.Background(), tc.amp)
 			if tc.expectError {