feat(image extender): set image for init container

go.mod

@@ -23,7 +23,7 @@ require (
 	github.com/ryanuber/go-glob v1.0.0 // indirect
 	github.com/sahilm/fuzzy v0.1.0
 	github.com/spf13/cobra v1.0.0
-	github.com/stretchr/testify v1.5.1
+	github.com/stretchr/testify v1.6.1
 	golang.org/x/net v0.0.0-20200519113804-d87ec0cfa476 // indirect
 	golang.org/x/sys v0.0.0-20200519105757-fe76b779f299 // indirect
 	golang.org/x/text v0.3.2

go.sum

@@ -615,6 +615,8 @@ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UV
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.5.1 h1:nOGnQDM7FYENwehXlg/kFVnos3rEvtKTjRvOWSzb6H4=
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
+github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
+github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/syndtr/gocapability v0.0.0-20170704070218-db04d3cc01c8/go.mod h1:hkRG7XYTFWNJGYcbNJQlaLq0fg1yr4J4t/NcTQtrfww=
 github.com/tidwall/pretty v1.0.0/go.mod h1:XNkn88O1ChpSDQmQeStsy+sBenx6DDtFZJxhVysOjyk=
 github.com/tmc/grpc-websocket-proxy v0.0.0-20170815181823-89b8d40f7ca8/go.mod h1:ncp9v5uamzpCO7NfCPTXjqaC+bZgJeR0sMTm6dMHP7U=
@@ -839,6 +841,8 @@ gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v2 v2.2.5/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10=
 gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gotest.tools v2.2.0+incompatible h1:VsBPFP1AI068pPrMxtb/S8Zkgf9xEmTLJjfM+P5UIEo=
 gotest.tools v2.2.0+incompatible/go.mod h1:DsYFclhRJ6vuDpmuTbkuFWG+y2sxOXAzmJt81HFBacw=
 helm.sh/helm/v3 v3.2.0 h1:V12EGAmr2DJ/fWrPo2fPdXWSIXvlXm51vGkQIXMeymE=

internal/dao/dp.go

@@ -220,7 +220,7 @@ func (d *Deployment) GetPodSpec(path string) (*v1.PodSpec, error) {
 	return &podSpec, nil
 }
 
-func (d *Deployment) SetImages(ctx context.Context, path string, images map[string]string) error {
+func (d *Deployment) SetImages(ctx context.Context, path string, spec v1.PodSpec) error {
 	ns, n := client.Namespaced(path)
 	auth, err := d.Client().CanI(ns, "apps/v1/deployments", []string{client.PatchVerb})
 	if err != nil {
@@ -229,8 +229,7 @@ func (d *Deployment) SetImages(ctx context.Context, path string, images map[stri
 	if !auth {
 		return fmt.Errorf("user is not authorized to patch a deployment")
 	}
-
+	jsonPatch, err := SetImageJsonPatch(spec)
-	jsonPatch, err := SetImageJsonPatch(images)
 	if err != nil {
 		return err
 	}

internal/dao/patch.go

@@ -1,6 +1,9 @@
 package dao
 
-import "encoding/json"
+import (
+	"encoding/json"
+	v1 "k8s.io/api/core/v1"
+)
 
 type JsonPatch struct {
 	Spec Spec `json:"spec"`
@@ -15,8 +18,10 @@ type Template struct {
 }
 
 type ImagesSpec struct {
-	SetElementOrders []Element `json:"$setElementOrder/containers"`
+	SetElementOrderContainers     []Element `json:"$setElementOrder/containers,omitempty"`
-	Containers       []Element `json:"containers"`
+	SetElementOrderInitContainers []Element `json:"$setElementOrder/initContainers,omitempty"`
+	Containers                    []Element `json:"containers,omitempty"`
+	InitContainers                []Element `json:"initContainers,omitempty"`
 }
 
 type Element struct {
@@ -25,19 +30,15 @@ type Element struct {
 }
 
 // Build a json patch string to update PodSpec images
-func SetImageJsonPatch(images map[string]string) (string, error) {
+func SetImageJsonPatch(spec v1.PodSpec) (string, error) {
-	elementOrders := make([]Element, 0)
-	containers := make([]Element, 0)
-	for key, value := range images {
-		elementOrders = append(elementOrders, Element{Name: key})
-		containers = append(containers, Element{Name: key, Image: value})
-	}
 	jsonPatch := JsonPatch{
 		Spec: Spec{
 			Template: Template{
 				Spec: ImagesSpec{
-					SetElementOrders: elementOrders,
+					SetElementOrderContainers:     extractElements(spec.Containers, false),
-					Containers:       containers,
+					Containers:                    extractElements(spec.Containers, true),
+					SetElementOrderInitContainers: extractElements(spec.InitContainers, false),
+					InitContainers:                extractElements(spec.InitContainers, true),
 				},
 			},
 		},
@@ -45,3 +46,15 @@ func SetImageJsonPatch(images map[string]string) (string, error) {
 	bytes, err := json.Marshal(jsonPatch)
 	return string(bytes), err
 }
+
+func extractElements(containers []v1.Container, withImage bool) []Element {
+	elements := make([]Element, 0)
+	for _, c := range containers {
+		if withImage {
+			elements = append(elements, Element{Name: c.Name, Image: c.Image})
+		} else {
+			elements = append(elements, Element{Name: c.Name})
+		}
+	}
+	return elements
+}

internal/dao/patch_test.go

@@ -1,13 +1,14 @@
 package dao
 
 import (
-	"reflect"
+	"github.com/stretchr/testify/require"
+	v1 "k8s.io/api/core/v1"
 	"testing"
 )
 
 func TestSetImageJsonPatch(t *testing.T) {
 	type args struct {
-		images map[string]string
+		podSpec v1.PodSpec
 	}
 	tests := []struct {
 		name    string
@@ -18,22 +19,23 @@ func TestSetImageJsonPatch(t *testing.T) {
 		{
 			name: "simple",
 			args: args{
-				images: map[string]string{"nginx": "nginx:latest"},
+				podSpec: v1.PodSpec{
+					InitContainers: []v1.Container{v1.Container{Image: "busybox:latest", Name: "init"}},
+					Containers:     []v1.Container{v1.Container{Image: "nginx:latest", Name: "nginx"}},
+				},
 			},
-			want:    "",
+			want:    `{"spec":{"template":{"spec":{"$setElementOrder/containers":[{"name":"nginx"}],"$setElementOrder/initContainers":[{"name":"init"}],"containers":[{"image":"nginx:latest","name":"nginx"}],"initContainers":[{"image":"busybox:latest","name":"init"}]}}}}`,
 			wantErr: false,
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := SetImageJsonPatch(tt.args.images)
+			got, err := SetImageJsonPatch(tt.args.podSpec)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("SetImageJsonPatch() error = %v, wantErr %v", err, tt.wantErr)
 				return
 			}
-			if !reflect.DeepEqual(got, tt.want) {
+			require.JSONEq(t, tt.want, got, "Json strings should be equal")
-				t.Errorf("SetImageJsonPatch() got = %v, want %v", got, tt.want)
-			}
 		})
 	}
 }

internal/dao/types.go

@@ -152,5 +152,5 @@ type ContainsPodSpec interface {
 	GetPodSpec(path string) (*v1.PodSpec, error)
 
 	// Set Images for a resource
-	SetImages(ctx context.Context, path string, images map[string]string) error
+	SetImages(ctx context.Context, path string, spec v1.PodSpec) error
 }

internal/view/set_image_extender.go

@@ -54,17 +54,29 @@ func (s *SetImageExtender) showSetImageDialog(path string) {
 	s.App().Content.ShowPage(setImageKey)
 }
 
+type ContainerType string
+
+var runningContainer = ContainerType("Container")
+var initContainer = ContainerType("InitContainer")
+
+type ContainerImage struct {
+	ContainerType ContainerType
+	Image         string
+}
+
 func (s *SetImageExtender) makeSetImageForm(sel string) *tview.Form {
 	f := s.makeStyledForm()
-	containers, err := s.getImages(sel)
+	podSpec, err := s.getPodSpec(sel)
+	originalImages := getImages(podSpec)
+	formSubmitResult := make(map[string]ContainerImage, 0)
 	if err != nil {
 		s.App().Flash().Err(err)
 		return nil
 	}
-	images := make(map[string]string)
+	for name, containerImage := range originalImages {
-	for _, container := range *containers {
+		f.AddInputField(name, containerImage.Image, 0, nil, func(changed string) {
-		f.AddInputField(container.Name, container.Image, 0, nil, func(changed string) {
+			log.Info().Msgf("changed : %v", changed)
-			images[container.Name] = changed
+			formSubmitResult[name] = ContainerImage{ContainerType: containerImage.ContainerType, Image: changed}
 		})
 	}
 
@@ -77,14 +89,15 @@ func (s *SetImageExtender) makeSetImageForm(sel string) *tview.Form {
 		}
 		ctx, cancel := context.WithTimeout(context.Background(), s.App().Conn().Config().CallTimeout())
 		defer cancel()
-		if err := s.setImages(ctx, sel, images); err != nil {
+		podSpecPatch := buildPodSpecPatch(formSubmitResult, originalImages)
+		if err := s.setImages(ctx, sel, podSpecPatch); err != nil {
+
 			log.Error().Err(err).Msgf("DP %s image update failed", sel)
 			s.App().Flash().Err(err)
 		} else {
 			s.App().Flash().Infof("Resource %s:%s image updated successfully", s.GVR(), sel)
 		}
 	})
-
 	f.AddButton("Cancel", func() {
 		s.dismissDialog()
 	})
@@ -92,6 +105,48 @@ func (s *SetImageExtender) makeSetImageForm(sel string) *tview.Form {
 	return f
 }
 
+func getImages(podSpec *corev1.PodSpec) map[string]ContainerImage {
+	results := make(map[string]ContainerImage, 0)
+	for _, c := range podSpec.Containers {
+		results[c.Name] = ContainerImage{
+			ContainerType: runningContainer,
+			Image:         c.Image,
+		}
+	}
+	for _, c := range podSpec.InitContainers {
+		results[c.Name] = ContainerImage{
+			ContainerType: initContainer,
+			Image:         c.Image,
+		}
+	}
+	return results
+}
+
+func buildPodSpecPatch(formImages map[string]ContainerImage, originalImages map[string]ContainerImage) corev1.PodSpec {
+	initContainers := make([]corev1.Container, 0)
+	containers := make([]corev1.Container, 0)
+	for name, containerImage := range formImages {
+		if originalImages[name].Image == containerImage.Image {
+			continue
+		}
+		container := corev1.Container{
+			Image: containerImage.Image,
+			Name:  name,
+		}
+		switch containerImage.ContainerType {
+		case runningContainer:
+			containers = append(containers, container)
+		case initContainer:
+			initContainers = append(initContainers, container)
+		}
+	}
+	result := corev1.PodSpec{
+		Containers:     containers,
+		InitContainers: initContainers,
+	}
+	return result
+}
+
 func (s *SetImageExtender) dismissDialog() {
 	s.App().Content.RemovePage(setImageKey)
 }
@@ -107,7 +162,7 @@ func (s *SetImageExtender) makeStyledForm() *tview.Form {
 	return f
 }
 
-func (s *SetImageExtender) getImages(path string) (*[]corev1.Container, error) {
+func (s *SetImageExtender) getPodSpec(path string) (*corev1.PodSpec, error) {
 	res, err := dao.AccessorFor(s.App().factory, s.GVR())
 	if err != nil {
 		return nil, err
@@ -117,11 +172,10 @@ func (s *SetImageExtender) getImages(path string) (*[]corev1.Container, error) {
 		return nil, fmt.Errorf("expecting a podSpecable resource for %q", s.GVR())
 	}
 	podSpec, err := podSpecable.GetPodSpec(path)
-	containers := podSpec.Containers
+	return podSpec, nil
-	return &containers, nil
 }
 
-func (s *SetImageExtender) setImages(ctx context.Context, path string, images map[string]string) error {
+func (s *SetImageExtender) setImages(ctx context.Context, path string, spec corev1.PodSpec) error {
 	res, err := dao.AccessorFor(s.App().factory, s.GVR())
 	if err != nil {
 		return err
@@ -132,5 +186,5 @@ func (s *SetImageExtender) setImages(ctx context.Context, path string, images ma
 		return fmt.Errorf("expecting a scalable resource for %q", s.GVR())
 	}
 
-	return deployment.SetImages(ctx, path, images)
+	return deployment.SetImages(ctx, path, spec)
 }