fix: args parse issue

Signed-off-by: yxxhero <aiopsclub@163.com>

pkg/argparser/args.go

@@ -17,6 +17,22 @@ type argMap struct {
 	flags []string
 }
 
+// isNewFlag checks if the given arg is a new flag
+func isNewFlag(flag string) bool {
+	return strings.HasPrefix(flag, "--") || strings.HasPrefix(flag, "-")
+}
+
+// removeEmptyArgs removes empty args from the given args
+func removeEmptyArgs(args []string) []string {
+	var newArgs []string
+	for _, arg := range args {
+		if len(arg) > 0 {
+			newArgs = append(newArgs, arg)
+		}
+	}
+	return newArgs
+}
+
 // SetArg sets a flag and value in the map
 func (a *argMap) SetArg(flag, arg string, isSpace bool) {
 	// if flag is empty, return
@@ -27,7 +43,7 @@ func (a *argMap) SetArg(flag, arg string, isSpace bool) {
 		keyarg := &keyVal{key: flag, val: arg, spaceFlag: isSpace}
 		a.m[flag] = append(a.m[flag], keyarg)
 		a.flags = append(a.flags, flag)
-	} else if flag == "--set" || flag == "-f" || flag == "--values" {
+	} else {
 		keyarg := &keyVal{key: flag, val: arg, spaceFlag: isSpace}
 		a.m[flag] = append(a.m[flag], keyarg)
 	}
@@ -38,54 +54,44 @@ func newArgMap() *argMap {
 	return &argMap{m: map[string][]*keyVal{}}
 }
 
-func GetArgs(args string, state *state.HelmState) []string {
+func analyzeArgs(am *argMap, args string) {
-	argsMap := newArgMap()
+	argsVals := removeEmptyArgs(strings.Split(args, " "))
-
-	if len(args) > 0 {
-		argsVals := strings.Split(args, " ")
 	for index, arg := range argsVals {
-			if strings.HasPrefix(arg, "--") {
+		if len(arg) == 0 {
+			continue
+		}
+		if isNewFlag(arg) {
 			argVal := strings.SplitN(arg, "=", 2)
 			if len(argVal) > 1 {
 				arg := argVal[0]
 				value := argVal[1]
-					argsMap.SetArg(arg, value, false)
+				am.SetArg(arg, value, false)
 			} else {
 				// check if next value is arg to flag
 				if index+1 < len(argsVals) {
 					nextVal := argsVals[index+1]
-						if strings.HasPrefix(nextVal, "--") {
+					if isNewFlag(nextVal) {
-							argsMap.SetArg(arg, "", false)
+						am.SetArg(arg, "", false)
 					} else {
-							argsMap.SetArg(arg, nextVal, true)
+						am.SetArg(arg, nextVal, true)
 					}
 				} else {
-						argsMap.SetArg(arg, "", false)
+					am.SetArg(arg, "", false)
 				}
 			}
 		}
 	}
 }
 
+func GetArgs(args string, state *state.HelmState) []string {
+	argsMap := newArgMap()
+
+	if len(args) > 0 {
+		analyzeArgs(argsMap, args)
+	}
+
 	if len(state.HelmDefaults.Args) > 0 {
-		for _, arg := range state.HelmDefaults.Args {
+		analyzeArgs(argsMap, strings.Join(state.HelmDefaults.Args, " "))
-			var flag string
-			var val string
-
-			argsNum, _ := fmt.Sscanf(arg, "--%s %s", &flag, &val)
-			if argsNum == 2 {
-				argsMap.SetArg(flag, val, true)
-			} else {
-				argVal := strings.SplitN(arg, "=", 2)
-				argFirst := argVal[0]
-				if len(argVal) > 1 {
-					val = argVal[1]
-					argsMap.SetArg(argFirst, val, false)
-				} else {
-					argsMap.SetArg(argFirst, "", false)
-				}
-			}
-		}
 	}
 
 	var argArr []string

pkg/argparser/args_test.go

@@ -14,19 +14,30 @@ func TestGetArgs(t *testing.T) {
 	tests := []struct {
 		args        string
 		expected    string
+		defaultArgs []string
 	}{
+		{
+			args:        "-f a.yaml -f b.yaml -i --set app1.bootstrap=true --set app2.bootstrap=false",
+			defaultArgs: []string{"--recreate-pods", "--force"},
+			expected:    "-f a.yaml -f b.yaml -i --set app1.bootstrap=true --set app2.bootstrap=false --recreate-pods --force",
+		},
+		{
+			args:        "-e  a.yaml   -d b.yaml   -i --set app1.bootstrap=true --set app2.bootstrap=false",
+			defaultArgs: []string{"-q www", "-w"},
+			expected:    "-e a.yaml -d b.yaml -i --set app1.bootstrap=true --set app2.bootstrap=false -q www -w",
+		},
 		{
 			args:     "--timeout=3600 --set app1.bootstrap=true --set app2.bootstrap=false --tiller-namespace ns",
-			expected: "--timeout=3600 --set app1.bootstrap=true --set app2.bootstrap=false --tiller-namespace ns --recreate-pods --force",
+			expected: "--timeout=3600 --set app1.bootstrap=true --set app2.bootstrap=false --tiller-namespace ns",
 		},
 		{
 			args:        "--timeout=3600 --set app1.bootstrap=true --set app2.bootstrap=false,app3.bootstrap=true --tiller-namespace ns",
+			defaultArgs: []string{"--recreate-pods", "--force"},
 			expected:    "--timeout=3600 --set app1.bootstrap=true --set app2.bootstrap=false,app3.bootstrap=true --tiller-namespace ns --recreate-pods --force",
 		},
 	}
 	for _, test := range tests {
-		defaultArgs := []string{"--recreate-pods", "--force"}
+		Helmdefaults := state.HelmSpec{KubeContext: "test", TillerNamespace: "test-namespace", Args: test.defaultArgs}
-		Helmdefaults := state.HelmSpec{KubeContext: "test", TillerNamespace: "test-namespace", Args: defaultArgs}
 		testState := &state.HelmState{
 			ReleaseSetSpec: state.ReleaseSetSpec{
 				HelmDefaults: Helmdefaults,
@@ -38,6 +49,41 @@ func TestGetArgs(t *testing.T) {
 	}
 }
 
+// TestIsNewFlag tests the isNewFlag function
+func TestIsNewFlag(t *testing.T) {
+	tests := []struct {
+		arg      string
+		expected bool
+	}{
+		{
+			arg:      "-f",
+			expected: true,
+		},
+		{
+			arg:      "--file",
+			expected: true,
+		},
+		{
+			arg:      "file",
+			expected: false,
+		},
+	}
+
+	for _, test := range tests {
+		received := isNewFlag(test.arg)
+		require.Equalf(t, test.expected, received, "isNewFlag expected %t, received %t", test.expected, received)
+	}
+}
+
+// TestRemoveEmptyArgs tests the removeEmptyArgs function
+func TestRemoveEmptyArgs(t *testing.T) {
+	args := []string{"", "--set", "", "app1.bootstrap=true", "", "--set", "", "app2.bootstrap=true", "", "--timeout", "", "3600", "", "--force", "", ""}
+	expectArgs := []string{"--set", "app1.bootstrap=true", "--set", "app2.bootstrap=true", "--timeout", "3600", "--force"}
+
+	receivedArgs := removeEmptyArgs(args)
+	require.Equalf(t, expectArgs, receivedArgs, "removeEmptyArgs expected %s, received %s", expectArgs, receivedArgs)
+}
+
 // TestSetArg tests the SetArg function
 func TestSetArg(t *testing.T) {
 	ap := newArgMap()
@@ -55,6 +101,12 @@ func TestSetArg(t *testing.T) {
 			isSpace: false,
 			change:  true,
 		},
+		{
+			flag:    "--set",
+			arg:     "app2.bootstrap=true",
+			isSpace: false,
+			change:  true,
+		},
 		{
 			flag:    "--timeout",
 			arg:     "3600",
@@ -88,3 +140,40 @@ func TestSetArg(t *testing.T) {
 		}
 	}
 }
+
+// TestAnalyzeArgs tests the analyzeArgs function
+func TestAnalyzeArgs(t *testing.T) {
+	ap := newArgMap()
+
+	tests := []struct {
+		arg     string
+		flag    string
+		val     string
+		isSpace bool
+	}{
+		{
+			arg:     "--set app1.bootstrap=true",
+			flag:    "--set",
+			isSpace: true,
+			val:     "app1.bootstrap=true",
+		},
+		{
+			arg:  "--set=app2.bootstrap=true",
+			flag: "--set",
+			val:  "app2.bootstrap=true",
+		},
+		{
+			arg:  "-f",
+			flag: "-f",
+			val:  "",
+		},
+	}
+
+	for _, test := range tests {
+		analyzeArgs(ap, test.arg)
+		require.Containsf(t, ap.flags, test.flag, "expected flag %s to be set", test.flag)
+		require.Containsf(t, ap.m, test.flag, "expected m %s to be set", test.flag)
+		kv := &keyVal{key: test.flag, val: test.val, spaceFlag: test.isSpace}
+		require.Containsf(t, ap.m[test.flag], kv, "expected %v in m[%s]", kv, test.flag)
+	}
+}