Merge pull request #101669 from carlory/fix-kubectl-1042

fix kubectl set env or resources not working for initcontainers

staging/src/k8s.io/kubectl/pkg/cmd/set/set_env.go

@@ -365,7 +365,9 @@ func (o *EnvOptions) RunEnv() error {
 	patches := CalculatePatches(infos, scheme.DefaultJSONEncoder(), func(obj runtime.Object) ([]byte, error) {
 		_, err := o.updatePodSpecForObject(obj, func(spec *v1.PodSpec) error {
 			resolutionErrorsEncountered := false
+			initContainers, _ := selectContainers(spec.InitContainers, o.ContainerSelector)
 			containers, _ := selectContainers(spec.Containers, o.ContainerSelector)
+			containers = append(containers, initContainers...)
 			objName, err := meta.NewAccessor().Name(obj)
 			if err != nil {
 				return err

staging/src/k8s.io/kubectl/pkg/cmd/set/set_env_test.go

@@ -666,3 +666,102 @@ func TestSetEnvFromResource(t *testing.T) {
 		})
 	}
 }
+
+func TestSetEnvRemoteWithSpecificContainers(t *testing.T) {
+	inputs := []struct {
+		name     string
+		args     []string
+		selector string
+
+		expectedContainers int
+	}{
+		{
+			name:               "all containers",
+			args:               []string{"deployments", "redis", "env=prod"},
+			selector:           "*",
+			expectedContainers: 2,
+		},
+		{
+			name:               "use wildcards to select some containers",
+			args:               []string{"deployments", "redis", "env=prod"},
+			selector:           "red*",
+			expectedContainers: 1,
+		},
+		{
+			name:               "single container",
+			args:               []string{"deployments", "redis", "env=prod"},
+			selector:           "redis",
+			expectedContainers: 1,
+		},
+	}
+
+	for _, input := range inputs {
+		mockDeployment := &appsv1.Deployment{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "redis",
+				Namespace: "test",
+			},
+			Spec: appsv1.DeploymentSpec{
+				Template: corev1.PodTemplateSpec{
+					Spec: corev1.PodSpec{
+						InitContainers: []corev1.Container{
+							{
+								Name:  "init",
+								Image: "redis",
+							},
+						},
+						Containers: []corev1.Container{
+							{
+								Name:  "redis",
+								Image: "redis",
+							},
+						},
+					},
+				},
+			},
+		}
+		t.Run(input.name, func(t *testing.T) {
+			tf := cmdtesting.NewTestFactory().WithNamespace("test")
+			defer tf.Cleanup()
+			tf.ClientConfigVal = &restclient.Config{ContentConfig: restclient.ContentConfig{GroupVersion: &schema.GroupVersion{Version: ""}}}
+			tf.Client = &fake.RESTClient{
+				GroupVersion:         schema.GroupVersion{Group: "", Version: "v1"},
+				NegotiatedSerializer: scheme.Codecs.WithoutConversion(),
+				Client: fake.CreateHTTPClient(func(req *http.Request) (*http.Response, error) {
+					switch p, m := req.URL.Path, req.Method; {
+					case p == "/namespaces/test/deployments/redis" && m == http.MethodGet:
+						return &http.Response{StatusCode: http.StatusOK, Header: cmdtesting.DefaultHeader(), Body: objBody(mockDeployment)}, nil
+					case p == "/namespaces/test/deployments/redis" && m == http.MethodPatch:
+						stream, err := req.GetBody()
+						if err != nil {
+							return nil, err
+						}
+						bytes, err := ioutil.ReadAll(stream)
+						if err != nil {
+							return nil, err
+						}
+						updated := strings.Count(string(bytes), `"value":`+`"`+"prod"+`"`)
+						if updated != input.expectedContainers {
+							t.Errorf("expected %d containers to be selected but got %d \n", input.expectedContainers, updated)
+						}
+						return &http.Response{StatusCode: http.StatusOK, Header: cmdtesting.DefaultHeader(), Body: objBody(mockDeployment)}, nil
+					default:
+						t.Errorf("%s: unexpected request: %#v\n%#v", input.name, req.URL, req)
+						return nil, nil
+					}
+				}),
+			}
+			streams := genericclioptions.NewTestIOStreamsDiscard()
+			opts := &EnvOptions{
+				PrintFlags:        genericclioptions.NewPrintFlags("").WithDefaultOutput("yaml").WithTypeSetter(scheme.Scheme),
+				ContainerSelector: input.selector,
+				Overwrite:         true,
+				IOStreams:         streams,
+			}
+			err := opts.Complete(tf, NewCmdEnv(tf, streams), input.args)
+			assert.NoError(t, err)
+			err = opts.RunEnv()
+			assert.NoError(t, err)
+		})
+	}
+}

staging/src/k8s.io/kubectl/pkg/cmd/set/set_resources.go

@@ -232,7 +232,9 @@ func (o *SetResourcesOptions) Run() error {
 	patches := CalculatePatches(o.Infos, scheme.DefaultJSONEncoder(), func(obj runtime.Object) ([]byte, error) {
 		transformed := false
 		_, err := o.UpdatePodSpecForObject(obj, func(spec *v1.PodSpec) error {
+			initContainers, _ := selectContainers(spec.InitContainers, o.ContainerSelector)
 			containers, _ := selectContainers(spec.Containers, o.ContainerSelector)
+			containers = append(containers, initContainers...)
 			if len(containers) != 0 {
 				for i := range containers {
 					if len(o.Limits) != 0 && len(containers[i].Resources.Limits) == 0 {

staging/src/k8s.io/kubectl/pkg/cmd/set/set_resources_test.go

@@ -516,3 +516,112 @@ func TestSetResourcesRemote(t *testing.T) {
 		})
 	}
 }
+
+func TestSetResourcesRemoteWithSpecificContainers(t *testing.T) {
+	inputs := []struct {
+		name     string
+		selector string
+		args     []string
+
+		expectedContainers int
+	}{
+		{
+			name:               "all containers",
+			args:               []string{"deployments", "redis"},
+			selector:           "*",
+			expectedContainers: 2,
+		},
+		{
+			name:               "use wildcards to select some containers",
+			args:               []string{"deployments", "redis"},
+			selector:           "red*",
+			expectedContainers: 1,
+		},
+		{
+			name:               "single container",
+			args:               []string{"deployments", "redis"},
+			selector:           "redis",
+			expectedContainers: 1,
+		},
+	}
+
+	for _, input := range inputs {
+		mockDeployment := &appsv1.Deployment{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "redis",
+				Namespace: "test",
+			},
+			Spec: appsv1.DeploymentSpec{
+				Template: corev1.PodTemplateSpec{
+					Spec: corev1.PodSpec{
+						InitContainers: []corev1.Container{
+							{
+								Name:  "init",
+								Image: "redis",
+							},
+						},
+						Containers: []corev1.Container{
+							{
+								Name:  "redis",
+								Image: "redis",
+							},
+						},
+					},
+				},
+			},
+		}
+		t.Run(input.name, func(t *testing.T) {
+			tf := cmdtesting.NewTestFactory().WithNamespace("test")
+			defer tf.Cleanup()
+			tf.ClientConfigVal = &restclient.Config{ContentConfig: restclient.ContentConfig{GroupVersion: &schema.GroupVersion{Version: ""}}}
+			tf.Client = &fake.RESTClient{
+				GroupVersion:         schema.GroupVersion{Group: "", Version: "v1"},
+				NegotiatedSerializer: scheme.Codecs.WithoutConversion(),
+				Client: fake.CreateHTTPClient(func(req *http.Request) (*http.Response, error) {
+					switch p, m := req.URL.Path, req.Method; {
+					case p == "/namespaces/test/deployments/redis" && m == http.MethodGet:
+						return &http.Response{StatusCode: http.StatusOK, Header: cmdtesting.DefaultHeader(), Body: objBody(mockDeployment)}, nil
+					case p == "/namespaces/test/deployments/redis" && m == http.MethodPatch:
+						stream, err := req.GetBody()
+						if err != nil {
+							return nil, err
+						}
+						bytes, err := ioutil.ReadAll(stream)
+						if err != nil {
+							return nil, err
+						}
+						updated := strings.Count(string(bytes), "200m")
+						if updated != input.expectedContainers {
+							t.Errorf("expected %d containers to be selected but got %d \n", input.expectedContainers, updated)
+						}
+						return &http.Response{StatusCode: http.StatusOK, Header: cmdtesting.DefaultHeader(), Body: objBody(mockDeployment)}, nil
+					default:
+						t.Errorf("%s: unexpected request: %#v\n%#v", input.name, req.URL, req)
+						return nil, nil
+					}
+				}),
+			}
+
+			outputFormat := "yaml"
+
+			streams := genericclioptions.NewTestIOStreamsDiscard()
+			cmd := NewCmdResources(tf, streams)
+			cmd.Flags().Set("output", outputFormat)
+			opts := SetResourcesOptions{
+				PrintFlags: genericclioptions.NewPrintFlags("").WithDefaultOutput(outputFormat).WithTypeSetter(scheme.Scheme),
+
+				Limits:            "cpu=200m,memory=512Mi",
+				ContainerSelector: input.selector,
+				IOStreams:         streams,
+			}
+			err := opts.Complete(tf, cmd, input.args)
+			if err == nil {
+				err = opts.Validate()
+			}
+			if err == nil {
+				err = opts.Run()
+			}
+			assert.NoError(t, err)
+		})
+	}
+}