From 9f7cad518207c6bfcc8b80d1b64681ffd91aa67c Mon Sep 17 00:00:00 2001
From: Jesse Suen <jessesuen@gmail.com>
Date: Tue, 12 Nov 2019 22:47:06 -0800
Subject: [PATCH] Add ability to `kubectl argo rollouts set image
 *=myrepo/myimage`

---
 .../cmd/set/set_image.go                      |  6 +-
 pkg/kubectl-argo-rollouts/cmd/set/set_test.go | 64 +++++++++++++++++++
 2 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/pkg/kubectl-argo-rollouts/cmd/set/set_image.go b/pkg/kubectl-argo-rollouts/cmd/set/set_image.go
index 717c009b72..d989cecba2 100644
--- a/pkg/kubectl-argo-rollouts/cmd/set/set_image.go
+++ b/pkg/kubectl-argo-rollouts/cmd/set/set_image.go
@@ -82,21 +82,21 @@ func newRolloutSetImage(orig *v1alpha1.Rollout, container string, image string)
 	ro := orig.DeepCopy()
 	containerFound := false
 	for i, ctr := range ro.Spec.Template.Spec.InitContainers {
-		if ctr.Name == container {
+		if ctr.Name == container || container == "*" {
 			containerFound = true
 			ctr.Image = image
 			ro.Spec.Template.Spec.InitContainers[i] = ctr
 		}
 	}
 	for i, ctr := range ro.Spec.Template.Spec.Containers {
-		if ctr.Name == container {
+		if ctr.Name == container || container == "*" {
 			containerFound = true
 			ctr.Image = image
 			ro.Spec.Template.Spec.Containers[i] = ctr
 		}
 	}
 	for i, ctr := range ro.Spec.Template.Spec.EphemeralContainers {
-		if ctr.Name == container {
+		if ctr.Name == container || container == "*" {
 			containerFound = true
 			ctr.Image = image
 			ro.Spec.Template.Spec.EphemeralContainers[i] = ctr
diff --git a/pkg/kubectl-argo-rollouts/cmd/set/set_test.go b/pkg/kubectl-argo-rollouts/cmd/set/set_test.go
index bb56ec2f98..f7228fb62a 100644
--- a/pkg/kubectl-argo-rollouts/cmd/set/set_test.go
+++ b/pkg/kubectl-argo-rollouts/cmd/set/set_test.go
@@ -120,6 +120,70 @@ func TestSetImageCmd(t *testing.T) {
 	assert.Empty(t, stderr)
 }
 
+func TestSetImageCmdStar(t *testing.T) {
+	ro := v1alpha1.Rollout{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "guestbook",
+			Namespace: metav1.NamespaceDefault,
+		},
+		Spec: v1alpha1.RolloutSpec{
+			Template: corev1.PodTemplateSpec{
+				Spec: corev1.PodSpec{
+					InitContainers: []corev1.Container{
+						{
+							Name:  "guestbook",
+							Image: "argoproj/rollouts-demo:blue",
+						},
+					},
+					Containers: []corev1.Container{
+						{
+							Name:  "foo",
+							Image: "alpine:3.8",
+						},
+						{
+							Name:  "guestbook",
+							Image: "argoproj/rollouts-demo:blue",
+						},
+						{
+							Name:  "bar",
+							Image: "alpine:3.8",
+						},
+					},
+					EphemeralContainers: []corev1.EphemeralContainer{
+						{
+							EphemeralContainerCommon: corev1.EphemeralContainerCommon{
+								Name:  "guestbook",
+								Image: "argoproj/rollouts-demo:blue",
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	tf, o := options.NewFakeArgoRolloutsOptions(&ro)
+	defer tf.Cleanup()
+
+	cmd := NewCmdSetImage(o)
+	cmd.PersistentPreRunE = o.PersistentPreRunE
+	cmd.SetArgs([]string{"guestbook", "*=argoproj/rollouts-demo:NEWIMAGE"})
+	err := cmd.Execute()
+	assert.Nil(t, err)
+
+	modifiedRo, err := o.RolloutsClientset().ArgoprojV1alpha1().Rollouts(metav1.NamespaceDefault).Get(ro.Name, metav1.GetOptions{})
+	assert.NoError(t, err)
+	assert.Equal(t, "argoproj/rollouts-demo:NEWIMAGE", modifiedRo.Spec.Template.Spec.Containers[1].Image)
+	assert.Equal(t, "argoproj/rollouts-demo:NEWIMAGE", modifiedRo.Spec.Template.Spec.Containers[0].Image)
+	assert.Equal(t, "argoproj/rollouts-demo:NEWIMAGE", modifiedRo.Spec.Template.Spec.Containers[2].Image)
+	assert.Equal(t, "argoproj/rollouts-demo:NEWIMAGE", modifiedRo.Spec.Template.Spec.InitContainers[0].Image)
+	assert.Equal(t, "argoproj/rollouts-demo:NEWIMAGE", modifiedRo.Spec.Template.Spec.EphemeralContainers[0].Image)
+
+	stdout := o.Out.(*bytes.Buffer).String()
+	stderr := o.ErrOut.(*bytes.Buffer).String()
+	assert.Equal(t, stdout, "rollout \"guestbook\" image updated\n")
+	assert.Empty(t, stderr)
+}
 func TestSetImageCmdRolloutNotFound(t *testing.T) {
 	tf, o := options.NewFakeArgoRolloutsOptions()
 	defer tf.Cleanup()