Merge pull request #77400 from Klaven/arg_validation

cmd/kubeadm/app/cmd/phases/join/BUILD

@@ -37,6 +37,7 @@ go_library(
         "//staging/src/k8s.io/client-go/util/cert:go_default_library",
         "//vendor/github.com/lithammer/dedent:go_default_library",
         "//vendor/github.com/pkg/errors:go_default_library",
+        "//vendor/github.com/spf13/cobra:go_default_library",
         "//vendor/k8s.io/klog:go_default_library",
         "//vendor/k8s.io/utils/exec:go_default_library",
     ],

cmd/kubeadm/app/cmd/phases/join/controlplanejoin.go

@@ -20,6 +20,8 @@ import (
 	"fmt"
 
 	"github.com/pkg/errors"
+	"github.com/spf13/cobra"
+
 	"k8s.io/kubernetes/cmd/kubeadm/app/cmd/options"
 	"k8s.io/kubernetes/cmd/kubeadm/app/cmd/phases/workflow"
 	kubeadmconstants "k8s.io/kubernetes/cmd/kubeadm/app/constants"
@@ -58,6 +60,7 @@ func NewControlPlaneJoinPhase() workflow.Phase {
 				Short:          "Join a machine as a control plane instance",
 				InheritFlags:   getControlPlaneJoinPhaseFlags("all"),
 				RunAllSiblings: true,
+				ArgsValidator:  cobra.NoArgs,
 			},
 			newEtcdLocalSubphase(),
 			newUpdateStatusSubphase(),
@@ -68,10 +71,11 @@ func NewControlPlaneJoinPhase() workflow.Phase {
 
 func newEtcdLocalSubphase() workflow.Phase {
 	return workflow.Phase{
-		Name:         "etcd",
-		Short:        "Add a new local etcd member",
-		Run:          runEtcdPhase,
-		InheritFlags: getControlPlaneJoinPhaseFlags("etcd"),
+		Name:          "etcd",
+		Short:         "Add a new local etcd member",
+		Run:           runEtcdPhase,
+		InheritFlags:  getControlPlaneJoinPhaseFlags("etcd"),
+		ArgsValidator: cobra.NoArgs,
 	}
 }
 
@@ -83,17 +87,19 @@ func newUpdateStatusSubphase() workflow.Phase {
 			kubeadmconstants.ClusterStatusConfigMapKey,
 			kubeadmconstants.KubeadmConfigConfigMap,
 		),
-		Run:          runUpdateStatusPhase,
-		InheritFlags: getControlPlaneJoinPhaseFlags("update-status"),
+		Run:           runUpdateStatusPhase,
+		InheritFlags:  getControlPlaneJoinPhaseFlags("update-status"),
+		ArgsValidator: cobra.NoArgs,
 	}
 }
 
 func newMarkControlPlaneSubphase() workflow.Phase {
 	return workflow.Phase{
-		Name:         "mark-control-plane",
-		Short:        "Mark a node as a control-plane",
-		Run:          runMarkControlPlanePhase,
-		InheritFlags: getControlPlaneJoinPhaseFlags("mark-control-plane"),
+		Name:          "mark-control-plane",
+		Short:         "Mark a node as a control-plane",
+		Run:           runMarkControlPlanePhase,
+		InheritFlags:  getControlPlaneJoinPhaseFlags("mark-control-plane"),
+		ArgsValidator: cobra.NoArgs,
 	}
 }
 

cmd/kubeadm/app/cmd/phases/join/controlplaneprepare.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 
 	"github.com/pkg/errors"
+	"github.com/spf13/cobra"
 
 	clientset "k8s.io/client-go/kubernetes"
 	"k8s.io/klog"
@@ -157,10 +158,11 @@ func newControlPlanePrepareKubeconfigSubphase() workflow.Phase {
 
 func newControlPlanePrepareControlPlaneSubphase() workflow.Phase {
 	return workflow.Phase{
-		Name:         "control-plane",
-		Short:        "Generate the manifests for the new control plane components",
-		Run:          runControlPlanePrepareControlPlaneSubphase, //NB. eventually in future we would like to break down this in sub phases for each component
-		InheritFlags: getControlPlanePreparePhaseFlags("control-plane"),
+		Name:          "control-plane",
+		Short:         "Generate the manifests for the new control plane components",
+		Run:           runControlPlanePrepareControlPlaneSubphase, //NB. eventually in future we would like to break down this in sub phases for each component
+		InheritFlags:  getControlPlanePreparePhaseFlags("control-plane"),
+		ArgsValidator: cobra.NoArgs,
 	}
 }
 

cmd/kubeadm/app/cmd/phases/workflow/phase.go

@@ -16,7 +16,10 @@ limitations under the License.
 
 package workflow
 
-import "github.com/spf13/pflag"
+import (
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+)
 
 // Phase provides an implementation of a workflow phase that allows
 // creation of new phases by simply instantiating a variable of this type.
@@ -71,6 +74,10 @@ type Phase struct {
 	// Nb. if two or phases have the same local flags, please consider using local flags in the parent command
 	// or additional flags defined in the phase runner.
 	LocalFlags *pflag.FlagSet
+
+	// ArgsValidator defines the positional arg function to be used for validating args for this phase
+	// If not set a phase will adopt the args of the top level command.
+	ArgsValidator cobra.PositionalArgs
 }
 
 // AppendPhase adds the given phase to the nested, ordered sequence of phases.

cmd/kubeadm/app/cmd/phases/workflow/runner.go

@@ -372,6 +372,12 @@ func (e *Runner) BindToCommand(cmd *cobra.Command) {
 		// if this phase has children (not a leaf) it doesn't accept any args
 		if len(p.Phases) > 0 {
 			phaseCmd.Args = cobra.NoArgs
+		} else {
+			if p.ArgsValidator == nil {
+				phaseCmd.Args = cmd.Args
+			} else {
+				phaseCmd.Args = p.ArgsValidator
+			}
 		}
 
 		// adds the command to parent

cmd/kubeadm/app/cmd/phases/workflow/runner_test.go

@@ -297,6 +297,134 @@ func phaseBuilder5(name string, flags *pflag.FlagSet) Phase {
 	}
 }
 
+type argTest struct {
+	args cobra.PositionalArgs
+	pass []string
+	fail []string
+}
+
+func phaseBuilder6(name string, args cobra.PositionalArgs, phases ...Phase) Phase {
+	return Phase{
+		Name:          name,
+		Short:         fmt.Sprintf("long description for %s ...", name),
+		Phases:        phases,
+		ArgsValidator: args,
+	}
+}
+
+// customArgs is a custom cobra.PositionArgs function
+func customArgs(cmd *cobra.Command, args []string) error {
+	for _, a := range args {
+		if a != "qux" {
+			return fmt.Errorf("arg %s does not equal qux", a)
+		}
+	}
+	return nil
+}
+
+func TestBindToCommandArgRequirements(t *testing.T) {
+
+	// because cobra.ExactArgs(1) == cobra.ExactArgs(3), it is needed
+	// to run test argument sets that both pass and fail to ensure the correct function was set.
+	var usecases = []struct {
+		name      string
+		runner    Runner
+		testCases map[string]argTest
+		cmd       *cobra.Command
+	}{
+		{
+			name: "leaf command, no defined args, follow parent",
+			runner: Runner{
+				Phases: []Phase{phaseBuilder("foo")},
+			},
+			testCases: map[string]argTest{
+				"phase foo": {
+					pass: []string{"one", "two", "three"},
+					fail: []string{"one", "two"},
+					args: cobra.ExactArgs(3),
+				},
+			},
+			cmd: &cobra.Command{
+				Use:  "init",
+				Args: cobra.ExactArgs(3),
+			},
+		},
+		{
+			name: "container cmd expect none, custom arg check for leaf",
+			runner: Runner{
+				Phases: []Phase{phaseBuilder6("foo", cobra.NoArgs,
+					phaseBuilder6("bar", cobra.ExactArgs(1)),
+					phaseBuilder6("baz", customArgs),
+				)},
+			},
+			testCases: map[string]argTest{
+				"phase foo": {
+					pass: []string{},
+					fail: []string{"one"},
+					args: cobra.NoArgs,
+				},
+				"phase foo bar": {
+					pass: []string{"one"},
+					fail: []string{"one", "two"},
+					args: cobra.ExactArgs(1),
+				},
+				"phase foo baz": {
+					pass: []string{"qux"},
+					fail: []string{"one"},
+					args: customArgs,
+				},
+			},
+			cmd: &cobra.Command{
+				Use:  "init",
+				Args: cobra.NoArgs,
+			},
+		},
+	}
+
+	for _, rt := range usecases {
+		t.Run(rt.name, func(t *testing.T) {
+
+			rt.runner.BindToCommand(rt.cmd)
+
+			// Checks that cmd gets a new phase subcommand
+			phaseCmd := getCmd(rt.cmd, "phase")
+			if phaseCmd == nil {
+				t.Error("cmd didn't have phase subcommand\n")
+				return
+			}
+
+			for c, args := range rt.testCases {
+
+				cCmd := getCmd(rt.cmd, c)
+				if cCmd == nil {
+					t.Errorf("cmd didn't have %s subcommand\n", c)
+					continue
+				}
+
+				// Ensure it is the expected function
+				if reflect.ValueOf(cCmd.Args).Pointer() != reflect.ValueOf(args.args).Pointer() {
+					t.Error("The function poiners where not equal.")
+				}
+
+				// Test passing argument set
+				err := cCmd.Args(cCmd, args.pass)
+
+				if err != nil {
+					t.Errorf("command %s should validate the args: %v\n %v", cCmd.Name(), args.pass, err)
+				}
+
+				// Test failing argument set
+				err = cCmd.Args(cCmd, args.fail)
+
+				if err == nil {
+					t.Errorf("command %s should fail to validate the args: %v\n %v", cCmd.Name(), args.pass, err)
+				}
+			}
+
+		})
+	}
+}
+
 func TestBindToCommand(t *testing.T) {
 
 	var dummy string