handle deamonSet in FixImages

kubernetes-sigs · Apr 12, 2021 · 3da05c0 · 3da05c0
1 parent 2b9be65
commit 3da05c0
Show file tree

Hide file tree

Showing 2 changed files with 107 additions and 31 deletions.
diff --git a/cmd/clusterctl/internal/util/objs.go b/cmd/clusterctl/internal/util/objs.go
@@ -19,6 +19,7 @@ package util
 import (
 	"github.com/pkg/errors"
 	appsv1 "k8s.io/api/apps/v1"
+	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	clusterctlv1 "sigs.k8s.io/cluster-api/cmd/clusterctl/api/v1alpha3"
 	"sigs.k8s.io/cluster-api/cmd/clusterctl/internal/scheme"
@@ -119,42 +120,83 @@ func FixImages(objs []unstructured.Unstructured, alterImageFunc func(image strin
 	// look for resources of kind Deployment and alter the image
 	for i := range objs {
 		o := &objs[i]
-		if o.GetKind() != deploymentKind {
-			continue
+		err := fixDeploymentImage(o, alterImageFunc)
+		if err != nil {
+			return nil, err
 		}
-
-		// Convert Unstructured into a typed object
-		d := &appsv1.Deployment{}
-		if err := scheme.Scheme.Convert(o, d, nil); err != nil {
+		err = fixDaemonSetImage(o, alterImageFunc)
+		if err != nil {
 			return nil, err
 		}
+		objs[i] = *o
+	}
+	return objs, nil
+}
 
-		// Alter the image
-		for j := range d.Spec.Template.Spec.Containers {
-			container := d.Spec.Template.Spec.Containers[j]
-			image, err := alterImageFunc(container.Image)
-			if err != nil {
-				return nil, errors.Wrapf(err, "failed to fix image for container %s in deployment %s", container.Name, d.Name)
-			}
-			container.Image = image
-			d.Spec.Template.Spec.Containers[j] = container
-		}
+func fixDeploymentImage(o *unstructured.Unstructured, alterImageFunc func(image string) (string, error)) error {
+	if o.GetKind() != deploymentKind {
+		return nil
+	}
 
-		for j := range d.Spec.Template.Spec.InitContainers {
-			container := d.Spec.Template.Spec.InitContainers[j]
-			image, err := alterImageFunc(container.Image)
-			if err != nil {
-				return nil, errors.Wrapf(err, "failed to fix image for init container %s in deployment %s", container.Name, d.Name)
-			}
-			container.Image = image
-			d.Spec.Template.Spec.InitContainers[j] = container
-		}
+	// Convert Unstructured into a typed object
+	d := &appsv1.Deployment{}
+	if err := scheme.Scheme.Convert(o, d, nil); err != nil {
+		return err
+	}
 
-		// Convert typed object back to Unstructured
-		if err := scheme.Scheme.Convert(d, o, nil); err != nil {
-			return nil, err
+	var err error
+	d.Spec.Template.Spec.Containers, err = fixContainersImage(d.Spec.Template.Spec.Containers, alterImageFunc)
+	if err != nil {
+		return errors.Wrapf(err, "failed to fix containers in deployment %s", d.Name)
+	}
+	d.Spec.Template.Spec.InitContainers, err = fixContainersImage(d.Spec.Template.Spec.InitContainers, alterImageFunc)
+	if err != nil {
+		return errors.Wrapf(err, "failed to fix init containers in deployment %s", d.Name)
+	}
+
+	// Convert typed object back to Unstructured
+	if err := scheme.Scheme.Convert(d, o, nil); err != nil {
+		return err
+	}
+	return nil
+}
+
+func fixDaemonSetImage(o *unstructured.Unstructured, alterImageFunc func(image string) (string, error)) error {
+	if o.GetKind() != daemonSetKind {
+		return nil
+	}
+
+	// Convert Unstructured into a typed object
+	d := &appsv1.DaemonSet{}
+	if err := scheme.Scheme.Convert(o, d, nil); err != nil {
+		return err
+	}
+
+	var err error
+	d.Spec.Template.Spec.Containers, err = fixContainersImage(d.Spec.Template.Spec.Containers, alterImageFunc)
+	if err != nil {
+		return errors.Wrapf(err, "failed to fix containers in deamonSet %s", d.Name)
+	}
+	d.Spec.Template.Spec.InitContainers, err = fixContainersImage(d.Spec.Template.Spec.InitContainers, alterImageFunc)
+	if err != nil {
+		return errors.Wrapf(err, "failed to fix init containers in deamonSet %s", d.Name)
+	}
+
+	// Convert typed object back to Unstructured
+	if err := scheme.Scheme.Convert(d, o, nil); err != nil {
+		return err
+	}
+	return nil
+}
+
+func fixContainersImage(containers []v1.Container, alterImageFunc func(image string) (string, error)) ([]v1.Container, error) {
+	for j := range containers {
+		container := &containers[j]
+		image, err := alterImageFunc(container.Image)
+		if err != nil {
+			return nil, errors.Wrapf(err, "failed to fix image for container %s", container.Name)
 		}
-		objs[i] = *o
+		container.Image = image
 	}
-	return objs, nil
+	return containers, nil
 }
diff --git a/cmd/clusterctl/internal/util/objs_test.go b/cmd/clusterctl/internal/util/objs_test.go
@@ -127,7 +127,7 @@ func Test_inspectImages(t *testing.T) {
 			wantErr: false,
 		},
 		{
-			name: "controller with init container",
+			name: "controller with deamonSet",
 			args: args{
 				objs: []unstructured.Unstructured{
 					{
@@ -221,6 +221,40 @@ func TestFixImages(t *testing.T) {
 			want:    []string{"foo-container-image", "foo-init-container-image"},
 			wantErr: false,
 		},
+		{
+			name: "fix daemonSet containers images",
+			args: args{
+				objs: []unstructured.Unstructured{
+					{
+						Object: map[string]interface{}{
+							"apiVersion": "apps/v1",
+							"kind":       daemonSetKind,
+							"spec": map[string]interface{}{
+								"template": map[string]interface{}{
+									"spec": map[string]interface{}{
+										"containers": []map[string]interface{}{
+											{
+												"image": "container-image",
+											},
+										},
+										"initContainers": []map[string]interface{}{
+											{
+												"image": "init-container-image",
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+				alterImageFunc: func(image string) (string, error) {
+					return fmt.Sprintf("foo-%s", image), nil
+				},
+			},
+			want:    []string{"foo-container-image", "foo-init-container-image"},
+			wantErr: false,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {