Kubectl taint node based on label selector

pkg/kubectl/cmd/taint.go

@@ -80,7 +80,10 @@ var (
 		kubectl taint nodes foo dedicated:NoSchedule-
 
 		# Remove from node 'foo' all the taints with key 'dedicated'
-		kubectl taint nodes foo dedicated-`))
+		kubectl taint nodes foo dedicated-
+
+		# Add a taint with key 'dedicated' on nodes having label mylabel=X
+		kubectl taint node -l myLabel=X  dedicated=foo:PreferNoSchedule`))
 )
 
 func NewCmdTaint(f cmdutil.Factory, out io.Writer) *cobra.Command {
@@ -243,27 +246,42 @@ func (o *TaintOptions) Complete(f cmdutil.Factory, out io.Writer, cmd *cobra.Com
 	if o.taintsToAdd, o.taintsToRemove, err = parseTaints(taintArgs); err != nil {
 		return cmdutil.UsageError(cmd, err.Error())
 	}
-
 	mapper, typer := f.Object()
 	o.builder = resource.NewBuilder(mapper, f.CategoryExpander(), typer, resource.ClientMapperFunc(f.ClientForMapping), f.Decoder(true)).
 		ContinueOnError().
 		NamespaceParam(namespace).DefaultNamespace()
+	if o.selector != "" {
+		o.builder = o.builder.SelectorParam(o.selector).ResourceTypes("node")
+	}
 	if o.all {
-		o.builder = o.builder.SelectAllParam(o.all).ResourceTypes("node")
-	} else {
-		if len(o.resources) < 2 {
-			return fmt.Errorf("at least one resource name must be specified since 'all' parameter is not set")
-		}
+		o.builder = o.builder.SelectAllParam(o.all).ResourceTypes("node").Flatten().Latest()
+	}
+	if !o.all && o.selector == "" && len(o.resources) >= 2 {
 		o.builder = o.builder.ResourceNames("node", o.resources[1:]...)
 	}
 	o.builder = o.builder.SelectorParam(o.selector).
 		Flatten().
 		Latest()
-
 	o.f = f
 	o.out = out
 	o.cmd = cmd
+	return nil
+}
 
+// validateFlags checks for the validation of flags for kubectl taints.
+func (o TaintOptions) validateFlags() error {
+	// Cannot have a non-empty selector and all flag set. They are mutually exclusive.
+	if o.all && o.selector != "" {
+		return fmt.Errorf("setting 'all' parameter with a non empty selector is prohibited.")
+	}
+	// If both selector and all are not set.
+	if !o.all && o.selector == "" {
+		if len(o.resources) < 2 {
+			return fmt.Errorf("at least one resource name must be specified since 'all' parameter is not set")
+		} else {
+			return nil
+		}
+	}
 	return nil
 }
 
@@ -297,8 +315,7 @@ func (o TaintOptions) Validate() error {
 	if len(conflictTaints) > 0 {
 		return fmt.Errorf("can not both modify and remove the following taint(s) in the same command: %s", strings.Join(conflictTaints, ", "))
 	}
-
-	return nil
+	return o.validateFlags()
 }
 
 // RunTaint does the work

pkg/kubectl/cmd/taint_test.go

@@ -86,6 +86,7 @@ func TestTaint(t *testing.T) {
 		args        []string
 		expectFatal bool
 		expectTaint bool
+		selector    bool
 	}{
 		// success cases
 		{
@@ -237,7 +238,6 @@ func TestTaint(t *testing.T) {
 
 	for _, test := range tests {
 		oldNode, expectNewNode := generateNodeAndTaintedNode(test.oldTaints, test.newTaints)
-
 		new_node := &v1.Node{}
 		tainted := false
 		f, tf, codec, ns := cmdtesting.NewAPIFactory()
@@ -248,6 +248,8 @@ func TestTaint(t *testing.T) {
 			Client: fake.CreateHTTPClient(func(req *http.Request) (*http.Response, error) {
 				m := &MyReq{req}
 				switch {
+				case m.isFor("GET", "/nodes"):
+					return &http.Response{StatusCode: 200, Header: defaultHeader(), Body: objBody(codec, oldNode)}, nil
 				case m.isFor("GET", "/nodes/node-name"):
 					return &http.Response{StatusCode: 200, Header: defaultHeader(), Body: objBody(codec, oldNode)}, nil
 				case m.isFor("PATCH", "/nodes/node-name"):
@@ -332,3 +334,50 @@ func TestTaint(t *testing.T) {
 		}
 	}
 }
+
+func TestValidateFlags(t *testing.T) {
+	tests := []struct {
+		taintOpts   TaintOptions
+		description string
+		expectFatal bool
+	}{
+
+		{
+			taintOpts:   TaintOptions{selector: "myLabel=X", all: false},
+			description: "With Selector and without All flag",
+			expectFatal: false,
+		},
+		{
+			taintOpts:   TaintOptions{selector: "", all: true},
+			description: "Without selector and All flag",
+			expectFatal: false,
+		},
+		{
+			taintOpts:   TaintOptions{selector: "myLabel=X", all: true},
+			description: "With Selector and with All flag",
+			expectFatal: true,
+		},
+		{
+			taintOpts:   TaintOptions{selector: "", all: false, resources: []string{"node"}},
+			description: "Without Selector and All flags and if node name is not provided",
+			expectFatal: true,
+		},
+		{
+			taintOpts:   TaintOptions{selector: "", all: false, resources: []string{"node", "node-name"}},
+			description: "Without Selector and ALL flags and if node name is provided",
+			expectFatal: false,
+		},
+	}
+	for _, test := range tests {
+		sawFatal := false
+		err := test.taintOpts.validateFlags()
+		if err != nil {
+			sawFatal = true
+		}
+		if test.expectFatal {
+			if !sawFatal {
+				t.Fatalf("%s expected not to fail", test.description)
+			}
+		}
+	}
+}