Chore: Refactor feature gate handling in workflow tests

Signed-off-by: vishal210893 <[email protected]>
Signed-off-by: semmet95 <[email protected]>

pkg/executor/workflow_test.go

@@ -19,8 +19,11 @@ package executor
 import (
 	"context"
 	"encoding/json"
+	"fmt"
+	"github.com/kubevela/workflow/pkg/features"
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
+	"k8s.io/component-base/featuregate"
 	"math"
-	"testing"
 	"time"
 
 	"cuelang.org/go/cue/cuecontext"
@@ -33,16 +36,13 @@ import (
 	kerrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
-	utilfeature "k8s.io/apiserver/pkg/util/feature"
-	featuregatetesting "k8s.io/component-base/featuregate/testing"
 	"k8s.io/utils/ptr"
 	"sigs.k8s.io/yaml"
 
 	monitorContext "github.com/kubevela/pkg/monitor/context"
 	"github.com/kubevela/workflow/api/v1alpha1"
 	wfContext "github.com/kubevela/workflow/pkg/context"
 	"github.com/kubevela/workflow/pkg/cue/process"
-	"github.com/kubevela/workflow/pkg/features"
 	"github.com/kubevela/workflow/pkg/providers"
 	"github.com/kubevela/workflow/pkg/tasks/builtin"
 	"github.com/kubevela/workflow/pkg/tasks/custom"
@@ -895,7 +895,8 @@ var _ = Describe("Test Workflow", func() {
 
 	It("Workflow test for failed after retries with suspend", func() {
 		By("Test failed-after-retries in StepByStep mode with suspend")
-		defer featuregatetesting.SetFeatureGateDuringTest(&testing.T{}, utilfeature.DefaultFeatureGate, features.EnableSuspendOnFailure, true)
+		cleanup := setFeatureGateForTest(utilfeature.DefaultFeatureGate, features.EnableSuspendOnFailure, true)
+		defer cleanup()
 		instance, runners := makeTestCase([]v1alpha1.WorkflowStep{
 			{
 				WorkflowStepBase: v1alpha1.WorkflowStepBase{
@@ -1469,7 +1470,8 @@ var _ = Describe("Test Workflow", func() {
 
 	It("Test failed after retries with sub steps", func() {
 		By("Test failed-after-retries with step group in StepByStep mode")
-		defer featuregatetesting.SetFeatureGateDuringTest(&testing.T{}, utilfeature.DefaultFeatureGate, features.EnableSuspendOnFailure, true)
+		cleanup := setFeatureGateForTest(utilfeature.DefaultFeatureGate, features.EnableSuspendOnFailure, true)
+		defer cleanup()
 		instance, runners := makeTestCase([]v1alpha1.WorkflowStep{
 			{
 				WorkflowStepBase: v1alpha1.WorkflowStepBase{
@@ -2441,6 +2443,30 @@ func cleanStepTimeStamp(wfStatus *v1alpha1.WorkflowRunStatus) {
 	}
 }
 
+func setFeatureGateForTest(gate featuregate.FeatureGate, f featuregate.Feature, value bool) func() {
+	// The Kubernetes feature gate interface requires a MutableVersionedFeatureGate for setting/resetting.
+	mutableGate, ok := gate.(featuregate.MutableVersionedFeatureGate)
+	Expect(ok).To(BeTrue(), "The provided feature gate must be a MutableVersionedFeatureGate to be modified in tests")
+
+	originalValue := mutableGate.Enabled(f)
+	originalExplicitlySet := mutableGate.ExplicitlySet(f)
+
+	// Set the desired value for the test.
+	err := mutableGate.Set(fmt.Sprintf("%s=%v", f, value))
+	Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("Failed to set feature gate %q to %v: %v", f, value, err))
+
+	// Return a cleanup function to restore the original state.
+	return func() {
+		if originalExplicitlySet {
+			err := mutableGate.Set(fmt.Sprintf("%s=%v", f, originalValue))
+			Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("Failed to restore feature gate %q to original value %v: %v", f, originalValue, err))
+		} else {
+			err := mutableGate.ResetFeatureValueToDefault(f)
+			Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("Failed to reset feature gate %q to default: %v", f, err))
+		}
+	}
+}
+
 const cmYaml = `apiVersion: v1
 data:
   components: '{"server":"{\"Scopes\":null,\"StandardWorkload\":\"{\\\"apiVersion\\\":\\\"v1\\\",\\\"kind\\\":\\\"Pod\\\",\\\"metadata\\\":{\\\"labels\\\":{\\\"app\\\":\\\"nginx\\\"}},\\\"spec\\\":{\\\"containers\\\":[{\\\"env\\\":[{\\\"name\\\":\\\"APP\\\",\\\"value\\\":\\\"nginx\\\"}],\\\"image\\\":\\\"nginx:1.14.2\\\",\\\"imagePullPolicy\\\":\\\"IfNotPresent\\\",\\\"name\\\":\\\"main\\\",\\\"ports\\\":[{\\\"containerPort\\\":8080,\\\"protocol\\\":\\\"TCP\\\"}]}]}}\",\"Traits\":[\"{\\\"apiVersion\\\":\\\"v1\\\",\\\"kind\\\":\\\"Service\\\",\\\"metadata\\\":{\\\"name\\\":\\\"my-service\\\"},\\\"spec\\\":{\\\"ports\\\":[{\\\"port\\\":80,\\\"protocol\\\":\\\"TCP\\\",\\\"targetPort\\\":8080}],\\\"selector\\\":{\\\"app\\\":\\\"nginx\\\"}}}\"]}"}'