Merge pull request #8219 from towca/jtuznik/drain-fix

Fix parsing --drain-priority-config

Author: jackfrancis

cluster-autoscaler/config/flags/flags.go

@@ -34,7 +34,7 @@ import (
 	"k8s.io/autoscaler/cluster-autoscaler/utils/units"
 
 	"k8s.io/client-go/rest"
-	klog "k8s.io/klog/v2"
+	"k8s.io/klog/v2"
 	kubelet_config "k8s.io/kubernetes/pkg/kubelet/apis/config"
 	scheduler_config "k8s.io/kubernetes/pkg/scheduler/apis/config"
 )
@@ -269,12 +269,12 @@ func createAutoscalingOptions() config.AutoscalingOptions {
 		klog.Fatalf("Failed to get scheduler config: %v", err)
 	}
 
-	if isFlagPassed("drain-priority-config") && isFlagPassed("max-graceful-termination-sec") {
+	if pflag.CommandLine.Changed("drain-priority-config") && pflag.CommandLine.Changed("max-graceful-termination-sec") {
 		klog.Fatalf("Invalid configuration, could not use --drain-priority-config together with --max-graceful-termination-sec")
 	}
 
 	var drainPriorityConfigMap []kubelet_config.ShutdownGracePeriodByPodPriority
-	if isFlagPassed("drain-priority-config") {
+	if pflag.CommandLine.Changed("drain-priority-config") {
 		drainPriorityConfigMap = parseShutdownGracePeriodsAndPriorities(*drainPriorityConfig)
 		if len(drainPriorityConfigMap) == 0 {
 			klog.Fatalf("Invalid configuration, parsing --drain-priority-config")
@@ -409,16 +409,6 @@ func createAutoscalingOptions() config.AutoscalingOptions {
 	}
 }
 
-func isFlagPassed(name string) bool {
-	found := false
-	flag.Visit(func(f *flag.Flag) {
-		if f.Name == name {
-			found = true
-		}
-	})
-	return found
-}
-
 func minMaxFlagString(min, max int64) string {
 	return fmt.Sprintf("%v:%v", min, max)
 }

cluster-autoscaler/config/flags/flags_test.go

@@ -17,11 +17,15 @@ limitations under the License.
 package flags
 
 import (
+	"flag"
 	"testing"
 
 	"k8s.io/autoscaler/cluster-autoscaler/config"
 	kubelet_config "k8s.io/kubernetes/pkg/kubelet/apis/config"
 
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	"github.com/spf13/pflag"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -146,3 +150,47 @@ func TestParseShutdownGracePeriodsAndPriorities(t *testing.T) {
 		})
 	}
 }
+
+func TestCreateAutoscalingOptions(t *testing.T) {
+	for _, tc := range []struct {
+		testName            string
+		flags               []string
+		wantOptionsAsserter func(t *testing.T, gotOptions config.AutoscalingOptions)
+	}{
+		{
+			testName: "DrainPriorityConfig defaults to an empty list when the flag isn't passed",
+			flags:    []string{},
+			wantOptionsAsserter: func(t *testing.T, gotOptions config.AutoscalingOptions) {
+				if diff := cmp.Diff([]kubelet_config.ShutdownGracePeriodByPodPriority{}, gotOptions.DrainPriorityConfig, cmpopts.EquateEmpty()); diff != "" {
+					t.Errorf("createAutoscalingOptions(): unexpected DrainPriorityConfig field (-want +got): %s", diff)
+				}
+			},
+		},
+		{
+			testName: "DrainPriorityConfig is parsed correctly when the flag passed",
+			flags:    []string{"--drain-priority-config", "5000:60,3000:50,0:40"},
+			wantOptionsAsserter: func(t *testing.T, gotOptions config.AutoscalingOptions) {
+				wantConfig := []kubelet_config.ShutdownGracePeriodByPodPriority{
+					{Priority: 5000, ShutdownGracePeriodSeconds: 60},
+					{Priority: 3000, ShutdownGracePeriodSeconds: 50},
+					{Priority: 0, ShutdownGracePeriodSeconds: 40},
+				}
+				if diff := cmp.Diff(wantConfig, gotOptions.DrainPriorityConfig); diff != "" {
+					t.Errorf("createAutoscalingOptions(): unexpected DrainPriorityConfig field (-want +got): %s", diff)
+				}
+			},
+		},
+	} {
+		t.Run(tc.testName, func(t *testing.T) {
+			pflag.CommandLine = pflag.NewFlagSet("test", pflag.ExitOnError)
+			pflag.CommandLine.AddGoFlagSet(flag.CommandLine)
+			err := pflag.CommandLine.Parse(tc.flags)
+			if err != nil {
+				t.Errorf("pflag.CommandLine.Parse() got unexpected error: %v", err)
+			}
+
+			gotOptions := createAutoscalingOptions()
+			tc.wantOptionsAsserter(t, gotOptions)
+		})
+	}
+}