Merge pull request #14 from DWSR/use-specified-container-image

coderanger · web-flow · commit 623b844e573a · 2022-06-14T23:49:50.000-07:00
fix: use image of container in MigratorSpec
diff --git a/webhook/webhook.go b/webhook/webhook.go
@@ -92,10 +92,18 @@ func (hook *initInjector) handleInner(ctx context.Context, req admission.Request
 
 	// For each migrator, inject an initContainer.
 	for _, m := range migrators {
+		// Look for the container named in the migrator and pull the image from that. If no container
+		// matches, fall back to using the first container.
+		podIdx := 0
+		for i, p := range pod.Spec.Containers {
+			if p.Name == m.Spec.Container {
+				podIdx = i
+			}
+		}
 		initContainer := corev1.Container{
 			Name:    fmt.Sprintf("migrate-wait-%s", m.Name),
 			Image:   os.Getenv("WAITER_IMAGE"),
-			Command: []string{"/waiter", pod.Spec.Containers[0].Image, m.Namespace, m.Name, os.Getenv("API_HOSTNAME")},
+			Command: []string{"/waiter", pod.Spec.Containers[podIdx].Image, m.Namespace, m.Name, os.Getenv("API_HOSTNAME")},
 			Resources: corev1.ResourceRequirements{
 				Requests: corev1.ResourceList{
 					corev1.ResourceMemory: resource.MustParse("16M"),
diff --git a/webhook/webhook_test.go b/webhook/webhook_test.go
@@ -94,6 +94,113 @@ var _ = Describe("InitInjector", func() {
 		Expect(pod.Spec.InitContainers[0].Command).To(Equal([]string{"/waiter", "fake", helper.Namespace, "testing", "migrations-operator.migration-operator.svc"}))
 	})
 
+	It("selects the specified container with a multi-container Pod", func() {
+		c := helper.TestClient
+
+		migrator := &migrationsv1beta1.Migrator{
+			ObjectMeta: metav1.ObjectMeta{Name: "testing"},
+			Spec: migrationsv1beta1.MigratorSpec{
+				Selector: &metav1.LabelSelector{
+					MatchLabels: map[string]string{"app": "testing"},
+				},
+				Container: "second",
+			},
+		}
+		c.Create(migrator)
+
+		pod := &corev1.Pod{
+			ObjectMeta: metav1.ObjectMeta{Name: "testing", Labels: map[string]string{"app": "testing"}},
+			Spec: corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name:  "main",
+						Image: "fake",
+					},
+					{
+						Name:  "second",
+						Image: "foo",
+					},
+				},
+			},
+		}
+		c.Create(pod)
+
+		c.EventuallyGetName("testing", pod)
+		Expect(pod.Spec.InitContainers).To(HaveLen(1))
+		Expect(pod.Spec.InitContainers[0].Command).To(Equal([]string{"/waiter", "foo", helper.Namespace, "testing", "migrations-operator.migration-operator.svc"}))
+	})
+
+	It("falls back to the first container if the specified container doesn't exist", func() {
+		c := helper.TestClient
+
+		migrator := &migrationsv1beta1.Migrator{
+			ObjectMeta: metav1.ObjectMeta{Name: "testing"},
+			Spec: migrationsv1beta1.MigratorSpec{
+				Selector: &metav1.LabelSelector{
+					MatchLabels: map[string]string{"app": "testing"},
+				},
+				Container: "third",
+			},
+		}
+		c.Create(migrator)
+
+		pod := &corev1.Pod{
+			ObjectMeta: metav1.ObjectMeta{Name: "testing", Labels: map[string]string{"app": "testing"}},
+			Spec: corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name:  "main",
+						Image: "bar",
+					},
+					{
+						Name:  "second",
+						Image: "foo",
+					},
+				},
+			},
+		}
+		c.Create(pod)
+
+		c.EventuallyGetName("testing", pod)
+		Expect(pod.Spec.InitContainers).To(HaveLen(1))
+		Expect(pod.Spec.InitContainers[0].Command).To(Equal([]string{"/waiter", "bar", helper.Namespace, "testing", "migrations-operator.migration-operator.svc"}))
+	})
+
+	It("uses the first container image if no container name is supplied with a multi-container Pod", func() {
+		c := helper.TestClient
+
+		migrator := &migrationsv1beta1.Migrator{
+			ObjectMeta: metav1.ObjectMeta{Name: "testing"},
+			Spec: migrationsv1beta1.MigratorSpec{
+				Selector: &metav1.LabelSelector{
+					MatchLabels: map[string]string{"app": "testing"},
+				},
+			},
+		}
+		c.Create(migrator)
+
+		pod := &corev1.Pod{
+			ObjectMeta: metav1.ObjectMeta{Name: "testing", Labels: map[string]string{"app": "testing"}},
+			Spec: corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name:  "main",
+						Image: "bar",
+					},
+					{
+						Name:  "second",
+						Image: "foo",
+					},
+				},
+			},
+		}
+		c.Create(pod)
+
+		c.EventuallyGetName("testing", pod)
+		Expect(pod.Spec.InitContainers).To(HaveLen(1))
+		Expect(pod.Spec.InitContainers[0].Command).To(Equal([]string{"/waiter", "bar", helper.Namespace, "testing", "migrations-operator.migration-operator.svc"}))
+	})
+
 	It("doesn't inject with a non-matching migrator", func() {
 		c := helper.TestClient
 
