From c752307a8545b9f1fd86e0895471e342fcc40124 Mon Sep 17 00:00:00 2001
From: Michael Aspinwall <maspinwall@google.com>
Date: Tue, 3 Feb 2026 21:49:09 +0000
Subject: [PATCH] Add SetForTesting for testing feature gates

Kubernetes-commit: 8099c570428806404b1a400143156c2ed0c67658
---
 features/envvar.go           | 20 ++++++++++++++++++--
 features/envvar_test.go      | 26 ++++++++++++++++++++++----
 features/testing/features.go | 18 ++++++++++++++++--
 3 files changed, 56 insertions(+), 8 deletions(-)

diff --git a/features/envvar.go b/features/envvar.go
index 8c3f887dc..d67c2adc9 100644
--- a/features/envvar.go
+++ b/features/envvar.go
@@ -121,12 +121,28 @@ func (f *envVarFeatureGates) Enabled(key Feature) bool {
 // Features set via this method take precedence over
 // the features set via environment variables.
 func (f *envVarFeatureGates) Set(featureName Feature, featureValue bool) error {
+	return f.set(featureName, featureValue, false)
+}
+
+// SetForTesting sets the given feature to the given value. This method
+// bypasses the check for locked features and should only be used for
+// testing purposes.
+//
+// Features set via this method take precedence over
+// the features set via environment variables.
+func (f *envVarFeatureGates) SetForTesting(featureName Feature, featureValue bool) error {
+	return f.set(featureName, featureValue, true)
+}
+
+func (f *envVarFeatureGates) set(featureName Feature, featureValue bool, allowChangingLockedFeatures bool) error {
 	feature, ok := f.known[featureName]
 	if !ok {
 		return fmt.Errorf("feature %q is not registered in FeatureGates %q", featureName, f.callSiteName)
 	}
-	if feature.LockToDefault && feature.Default != featureValue {
-		return fmt.Errorf("cannot set feature gate %q to %v, feature is locked to %v", featureName, featureValue, feature.Default)
+	if !allowChangingLockedFeatures {
+		if feature.LockToDefault && feature.Default != featureValue {
+			return fmt.Errorf("cannot set feature gate %q to %v, feature is locked to %v", featureName, featureValue, feature.Default)
+		}
 	}
 
 	f.lockEnabledViaSetMethod.Lock()
diff --git a/features/envvar_test.go b/features/envvar_test.go
index b032d06cf..e640af4e4 100644
--- a/features/envvar_test.go
+++ b/features/envvar_test.go
@@ -47,10 +47,11 @@ func TestEnvVarFeatureGates(t *testing.T) {
 	}
 
 	scenarios := []struct {
-		name              string
-		features          map[Feature]FeatureSpec
-		envVariables      map[string]string
-		setMethodFeatures map[Feature]bool
+		name                  string
+		features              map[Feature]FeatureSpec
+		envVariables          map[string]string
+		setMethodFeatures     map[Feature]bool
+		setForTestingFeatures map[Feature]bool
 
 		expectedFeaturesState                           map[Feature]bool
 		expectedInternalEnabledViaEnvVarFeatureState    map[Feature]bool
@@ -134,6 +135,19 @@ func TestEnvVarFeatureGates(t *testing.T) {
 			expectedInternalEnabledViaEnvVarFeatureState:    map[Feature]bool{"TestAlpha": true},
 			expectedInternalEnabledViaSetMethodFeatureState: map[Feature]bool{"TestAlpha": false},
 		},
+		{
+			name: "setting a feature via the SetForTesting method works even when it is locked",
+			features: map[Feature]FeatureSpec{
+				"TestAlpha": {
+					Default:       true,
+					LockToDefault: true,
+					PreRelease:    "Alpha",
+				},
+			},
+			setForTestingFeatures:                           map[Feature]bool{"TestAlpha": false},
+			expectedFeaturesState:                           map[Feature]bool{"TestAlpha": false},
+			expectedInternalEnabledViaSetMethodFeatureState: map[Feature]bool{"TestAlpha": false},
+		},
 	}
 	for _, scenario := range scenarios {
 		t.Run(scenario.name, func(t *testing.T) {
@@ -146,6 +160,10 @@ func TestEnvVarFeatureGates(t *testing.T) {
 				err := target.Set(k, v)
 				require.NoError(t, err)
 			}
+			for k, v := range scenario.setForTestingFeatures {
+				err := target.SetForTesting(k, v)
+				require.NoError(t, err)
+			}
 			for expectedFeature, expectedValue := range scenario.expectedFeaturesState {
 				actualValue := target.Enabled(expectedFeature)
 				require.Equal(t, expectedValue, actualValue, "expected feature=%v, to be=%v, not=%v", expectedFeature, expectedValue, actualValue)
diff --git a/features/testing/features.go b/features/testing/features.go
index b5555ff80..b01d63d6a 100644
--- a/features/testing/features.go
+++ b/features/testing/features.go
@@ -40,6 +40,12 @@ type featureGatesSetter interface {
 	Set(clientfeatures.Feature, bool) error
 }
 
+type featureGatesSetterForTesting interface {
+	featureGatesSetter
+
+	SetForTesting(clientfeatures.Feature, bool) error
+}
+
 // SetFeatureDuringTest sets the specified feature to the specified value for the duration of the test.
 //
 // Example use:
@@ -68,9 +74,17 @@ func setFeatureDuringTestInternal(tb testing.TB, feature clientfeatures.Feature,
 		}
 	}
 
-	if err := featureGates.Set(feature, featureValue); err != nil {
+	setFeature := func(feature clientfeatures.Feature, featureValue bool) error {
+		if testSetter, ok := featureGates.(featureGatesSetterForTesting); ok {
+			return testSetter.SetForTesting(feature, featureValue)
+		}
+		return featureGates.Set(feature, featureValue)
+	}
+
+	if err := setFeature(feature, featureValue); err != nil {
 		return err
 	}
+
 	overriddenFeatures[feature] = tb.Name()
 
 	tb.Cleanup(func() {
@@ -78,7 +92,7 @@ func setFeatureDuringTestInternal(tb testing.TB, feature clientfeatures.Feature,
 		defer overriddenFeaturesLock.Unlock()
 		delete(overriddenFeatures, feature)
 		// if default is not set
-		if err := featureGates.Set(feature, originalFeatureValue); err != nil {
+		if err := setFeature(feature, originalFeatureValue); err != nil {
 			tb.Errorf("failed restoring client-go feature: %v to its original value: %v, err: %v", feature, originalFeatureValue, err)
 		}
 	})