cleanup: make the featuragate test more extensible.

Signed-off-by: Lan Liang <[email protected]>

Author: liangyuanpeng

server/embed/config_test.go

@@ -92,136 +92,99 @@ func TestConfigFileOtherFields(t *testing.T) {
 	assert.Equal(t, false, cfg.SocketOpts.ReuseAddress, "ReuseAddress does not match")
 }
 
-func TestConfigFileFeatureGates(t *testing.T) {
-	testCases := []struct {
-		name                                string
-		serverFeatureGatesJSON              string
-		experimentalStopGRPCServiceOnDefrag string
-		experimentalInitialCorruptCheck     string
-		expectErr                           bool
-		expectedFeatures                    map[featuregate.Feature]bool
-	}{
-		{
-			name: "default",
-			expectedFeatures: map[featuregate.Feature]bool{
-				features.DistributedTracing:      false,
-				features.StopGRPCServiceOnDefrag: false,
-				features.InitialCorruptCheck:     false,
-			},
-		},
+func defaultFeatureGates() map[featuregate.Feature]bool {
+	fgs := make(map[featuregate.Feature]bool)
+	for fg, fgspec := range features.DefaultEtcdServerFeatureGates {
+		fgs[fg] = fgspec.Default
+	}
+	return fgs
+}
+
+func defaultFeatureGatesWithFG(fg featuregate.Feature, enabled bool) map[featuregate.Feature]bool {
+	fgs := defaultFeatureGates()
+	fgs[fg] = enabled
+	return fgs
+}
+
+func defaultFeatureGatesWithFGs(fgs2 []featuregate.Feature, enabled bool) map[featuregate.Feature]bool {
+	fgs := defaultFeatureGates()
+	for _, f := range fgs2 {
+		fgs[f] = enabled
+	}
+	return fgs
+}
+
+func testCasesForFeaturegate(experimentalKey, fgname string, fg featuregate.Feature) []featureTestCase {
+	return []featureTestCase{
 		{
-			name:                                "cannot set both experimental flag and feature gate flag",
-			serverFeatureGatesJSON:              "StopGRPCServiceOnDefrag=true",
-			experimentalStopGRPCServiceOnDefrag: "false",
-			expectErr:                           true,
+			name:                   "cannot set both experimental flag and feature gate flag",
+			serverFeatureGatesJSON: fgname + "=true",
+			experimentalKey:        experimentalKey,
+			experimentalValue:      false,
+			expectErr:              true,
 		},
 		{
-			name:                            "cannot set both experimental flag and feature gate flag for InitialCorruptCheck",
-			serverFeatureGatesJSON:          "InitialCorruptCheck=true",
-			experimentalInitialCorruptCheck: "false",
-			expectErr:                       true,
+			name:              "can set feature gate to true from experimental flag",
+			experimentalKey:   experimentalKey,
+			experimentalValue: true,
+			expectedFeatures:  defaultFeatureGatesWithFG(fg, true),
 		},
 		{
-			name:                                "ok to set different experimental flag and feature gate flag",
-			serverFeatureGatesJSON:              "DistributedTracing=true",
-			experimentalStopGRPCServiceOnDefrag: "true",
-			expectedFeatures: map[featuregate.Feature]bool{
-				features.DistributedTracing:      true,
-				features.StopGRPCServiceOnDefrag: true,
-				features.InitialCorruptCheck:     false,
-			},
+			name:              "can set feature gate to false from experimental flag",
+			experimentalKey:   experimentalKey,
+			experimentalValue: false,
+			expectedFeatures:  defaultFeatureGatesWithFG(fg, false),
 		},
 		{
-			name:                                "can set feature gate to true from experimental flag",
-			experimentalStopGRPCServiceOnDefrag: "true",
-			expectedFeatures: map[featuregate.Feature]bool{
-				features.StopGRPCServiceOnDefrag: true,
-				features.DistributedTracing:      false,
-				features.InitialCorruptCheck:     false,
-			},
+			name:                   "can set feature gate " + fgname + " to true from feature gate flag",
+			serverFeatureGatesJSON: fgname + "=true",
+			expectedFeatures:       defaultFeatureGatesWithFG(fg, true),
 		},
 		{
-			name:                                "can set feature gate to false from experimental flag",
-			experimentalStopGRPCServiceOnDefrag: "false",
-			expectedFeatures: map[featuregate.Feature]bool{
-				features.StopGRPCServiceOnDefrag: false,
-				features.DistributedTracing:      false,
-				features.InitialCorruptCheck:     false,
-			},
-		},
-		{
-			name:                            "can set feature gate experimentalInitialCorruptCheck to true from experimental flag",
-			experimentalInitialCorruptCheck: "true",
-			expectedFeatures: map[featuregate.Feature]bool{
-				features.StopGRPCServiceOnDefrag: false,
-				features.DistributedTracing:      false,
-				features.InitialCorruptCheck:     true,
-			},
-		},
-		{
-			name:                            "can set feature gate experimentalInitialCorruptCheck to false from experimental flag",
-			experimentalInitialCorruptCheck: "false",
-			expectedFeatures: map[featuregate.Feature]bool{
-				features.StopGRPCServiceOnDefrag: false,
-				features.DistributedTracing:      false,
-				features.InitialCorruptCheck:     false,
-			},
-		},
-		{
-			name:                   "can set feature gate StopGRPCServiceOnDefrag to true from feature gate flag",
-			serverFeatureGatesJSON: "StopGRPCServiceOnDefrag=true",
-			expectedFeatures: map[featuregate.Feature]bool{
-				features.StopGRPCServiceOnDefrag: true,
-				features.DistributedTracing:      false,
-				features.InitialCorruptCheck:     false,
-			},
+			name:                   "can set feature gate to false from feature gate flag",
+			serverFeatureGatesJSON: fgname + "=false",
+			expectedFeatures:       defaultFeatureGatesWithFG(fg, false),
 		},
+	}
+}
+
+type featureTestCase struct {
+	name                   string
+	serverFeatureGatesJSON string
+	expectErr              bool
+	expectedFeatures       map[featuregate.Feature]bool
+	experimentalKey        string
+	experimentalValue      bool
+}
+
+func TestConfigFileFeatureGates(t *testing.T) {
+	testCases := []featureTestCase{
 		{
-			name:                   "can set feature gate InitialCorruptCheck to true from feature gate flag",
-			serverFeatureGatesJSON: "InitialCorruptCheck=true",
-			expectedFeatures: map[featuregate.Feature]bool{
-				features.StopGRPCServiceOnDefrag: false,
-				features.DistributedTracing:      false,
-				features.InitialCorruptCheck:     true,
-			},
+			name:             "default",
+			expectedFeatures: defaultFeatureGates(),
 		},
 		{
-			name:                   "can set feature gate to false from feature gate flag",
-			serverFeatureGatesJSON: "StopGRPCServiceOnDefrag=false",
-			expectedFeatures: map[featuregate.Feature]bool{
-				features.StopGRPCServiceOnDefrag: false,
-				features.DistributedTracing:      false,
-				features.InitialCorruptCheck:     false,
-			},
+			name:                   "ok to set different experimental flag and feature gate flag",
+			serverFeatureGatesJSON: "DistributedTracing=true",
+			experimentalKey:        "experimental-stop-grpc-service-on-defrag",
+			experimentalValue:      true,
+			expectedFeatures:       defaultFeatureGatesWithFGs([]featuregate.Feature{features.StopGRPCServiceOnDefrag, features.DistributedTracing}, true),
 		},
 	}
+
+	testCases = append(testCases, testCasesForFeaturegate("experimental-stop-grpc-service-on-defrag", "StopGRPCServiceOnDefrag", features.StopGRPCServiceOnDefrag)...)
+	testCases = append(testCases, testCasesForFeaturegate("experimental-initial-corrupt-check", "InitialCorruptCheck", features.InitialCorruptCheck)...)
+
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			yc := struct {
-				ExperimentalStopGRPCServiceOnDefrag *bool  `json:"experimental-stop-grpc-service-on-defrag,omitempty"`
-				ExperimentalInitialCorruptCheck     *bool  `json:"experimental-initial-corrupt-check,omitempty"`
-				ServerFeatureGatesJSON              string `json:"feature-gates"`
-			}{
-				ServerFeatureGatesJSON: tc.serverFeatureGatesJSON,
-			}
+			ycmap := make(map[string]interface{})
+			ycmap["feature-gates"] = tc.serverFeatureGatesJSON
 
-			if tc.experimentalInitialCorruptCheck != "" {
-				experimentalInitialCorruptCheck, err := strconv.ParseBool(tc.experimentalInitialCorruptCheck)
-				if err != nil {
-					t.Fatal(err)
-				}
-				yc.ExperimentalInitialCorruptCheck = &experimentalInitialCorruptCheck
-			}
-
-			if tc.experimentalStopGRPCServiceOnDefrag != "" {
-				experimentalStopGRPCServiceOnDefrag, err := strconv.ParseBool(tc.experimentalStopGRPCServiceOnDefrag)
-				if err != nil {
-					t.Fatal(err)
-				}
-				yc.ExperimentalStopGRPCServiceOnDefrag = &experimentalStopGRPCServiceOnDefrag
+			if tc.experimentalKey != "" {
+				ycmap[tc.experimentalKey] = tc.experimentalValue
 			}
 
-			b, err := yaml.Marshal(&yc)
+			b, err := yaml.Marshal(ycmap)
 			if err != nil {
 				t.Fatal(err)
 			}