PIPE-1664: job level image attribute take precedence over controller podSpecPatch

internal/controller/scheduler/build_inputs.go

@@ -0,0 +1,22 @@
+package scheduler
+
+import v1 "k8s.io/api/core/v1"
+
+func applyCustomImageIfPresent(podSpec *v1.PodSpec, inputs *buildInputs) *v1.PodSpec {
+
+	customImage := inputs.envMap["BUILDKITE_IMAGE"]
+	if customImage == "" {
+		return podSpec
+	}
+
+	for i := range podSpec.Containers {
+		c := &podSpec.Containers[i]
+		if c.Name != CommandContainerName {
+			continue
+		}
+
+		c.Image = customImage
+	}
+
+	return podSpec
+}

internal/controller/scheduler/scheduler.go

@@ -43,6 +43,7 @@ const (
 	CopyAgentContainerName        = "copy-agent"
 	ImageCheckContainerNamePrefix = "imagecheck-"
 	CheckoutContainerName         = "checkout"
+	CommandContainerName          = "container-0"
 )
 
 var errK8sPluginProhibited = errors.New("the kubernetes plugin is prohibited by this controller, but was configured on this job")
@@ -419,14 +420,10 @@ func (w *worker) Build(podSpec *corev1.PodSpec, skipCheckout bool, inputs buildI
 	}
 
 	if len(podSpec.Containers) == 0 {
-		image := w.cfg.Image
-		if customImage := inputs.envMap["BUILDKITE_IMAGE"]; customImage != "" {
-			image = inputs.envMap["BUILDKITE_IMAGE"]
-		}
-		// Create a default command container named "container-0".
+		// Create a default command container
 		c := corev1.Container{
-			Name:         "container-0",
-			Image:        image,
+			Name:         CommandContainerName,
+			Image:        w.cfg.Image,
 			Command:      commandContainerCommand,
 			Args:         commandContainerArgs,
 			WorkingDir:   "/workspace",
@@ -446,7 +443,7 @@ func (w *worker) Build(podSpec *corev1.PodSpec, skipCheckout bool, inputs buildI
 				},
 				{
 					Name:  "BUILDKITE_SOCKETS_PATH",
-					Value: "/workspace/sockets/container-0",
+					Value: "/workspace/sockets/" + CommandContainerName,
 				},
 			},
 		}
@@ -682,6 +679,9 @@ func (w *worker) Build(podSpec *corev1.PodSpec, skipCheckout bool, inputs buildI
 		w.logger.Debug("Applied podSpec patch from controller", zap.Any("patched", patched))
 	}
 
+	// Support `image: ` syntax, this HAS TO happen between controller podSpec patch and plugin podSpec patch.
+	podSpec = applyCustomImageIfPresent(podSpec, &inputs)
+
 	// If present, patch from the k8s plugin is applied second.
 	if inputs.k8sPlugin != nil && inputs.k8sPlugin.PodSpecPatch != nil {
 		patched, err := PatchPodSpec(podSpec, inputs.k8sPlugin.PodSpecPatch, w.cfg.DefaultCommandParams, inputs.k8sPlugin, w.cfg.AllowPodSpecPatchUnsafeCmdMod)

internal/controller/scheduler/scheduler_test.go

@@ -161,6 +161,39 @@ func TestPatchPodSpec(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "can patch image",
+			podspec: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name:  "container-0",
+						Image: "alpine:a",
+						Command: []string{
+							"echo hello world",
+						},
+					},
+				},
+			},
+			patch: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name:  "container-0",
+						Image: "alpine:b",
+					},
+				},
+			},
+			want: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name:  "container-0",
+						Image: "alpine:b",
+						Command: []string{
+							"echo hello world",
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for _, test := range cases {
@@ -882,6 +915,73 @@ func TestProhibitKubernetesPlugin(t *testing.T) {
 	require.Error(t, err)
 }
 
+func TestCustomImageSyntax_pluginTakesTopPriority(t *testing.T) {
+	t.Parallel()
+
+	pluginsYAML := `- github.com/buildkite-plugins/kubernetes-buildkite-plugin:
+    podSpecPatch:
+      containers:
+      - name: container-0
+        image: "x-image:plugin"`
+
+	pluginsJSON, err := yaml.YAMLToJSONStrict([]byte(pluginsYAML))
+	require.NoError(t, err)
+
+	job := &api.AgentJob{
+		ID: "abc",
+		Env: map[string]string{
+			"BUILDKITE_PLUGINS": string(pluginsJSON),
+			"BUILDKITE_IMAGE":   "x-image:job",
+		},
+	}
+	sjob := &api.AgentScheduledJob{
+		AgentQueryRules: []string{"queue=kubernetes"},
+	}
+	worker := New(zaptest.NewLogger(t), nil, nil, Config{
+		Image: "buildkite/agent:latest",
+	})
+	inputs, err := worker.ParseJob(job, sjob)
+	require.NoError(t, err)
+	kjob, err := worker.Build(&corev1.PodSpec{}, false, inputs)
+	require.NoError(t, err)
+
+	commandContainer := findContainer(t, kjob.Spec.Template.Spec.Containers, CommandContainerName)
+	require.Equal(t, "x-image:plugin", commandContainer.Image)
+}
+
+// Job level image syntax takes priority over controller setting
+func TestCustomImageSyntax_jobLevelImagePriority(t *testing.T) {
+	t.Parallel()
+
+	job := &api.AgentJob{
+		ID: "abc",
+		Env: map[string]string{
+			"BUILDKITE_IMAGE": "x-image:job",
+		},
+	}
+	sjob := &api.AgentScheduledJob{
+		AgentQueryRules: []string{"queue=kubernetes"},
+	}
+	worker := New(zaptest.NewLogger(t), nil, nil, Config{
+		Image: "buildkite/agent:latest",
+		PodSpecPatch: &corev1.PodSpec{
+			Containers: []corev1.Container{
+				{
+					Name:  "container-0",
+					Image: "alpine:controller",
+				},
+			},
+		},
+	})
+	inputs, err := worker.ParseJob(job, sjob)
+	require.NoError(t, err)
+	kjob, err := worker.Build(&corev1.PodSpec{}, false, inputs)
+	require.NoError(t, err)
+
+	commandContainer := findContainer(t, kjob.Spec.Template.Spec.Containers, CommandContainerName)
+	require.Equal(t, "x-image:job", commandContainer.Image)
+}
+
 func TestImagePullPolicies(t *testing.T) {
 	t.Parallel()