From 037091284e21741bb0427c489694e84774fff8cc Mon Sep 17 00:00:00 2001
From: Daniel Vega-Myhre <danielvm@google.com>
Date: Thu, 13 Jul 2023 22:38:21 +0000
Subject: [PATCH] fix unit test bug

---
 pkg/controller/job/indexed_job_utils.go   | 10 +++++++++-
 pkg/controller/job/job_controller_test.go | 12 +++++++++---
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/pkg/controller/job/indexed_job_utils.go b/pkg/controller/job/indexed_job_utils.go
index ef6b4e23426..43822a95e9a 100644
--- a/pkg/controller/job/indexed_job_utils.go
+++ b/pkg/controller/job/indexed_job_utils.go
@@ -26,8 +26,10 @@ import (
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/apiserver/pkg/storage/names"
+	"k8s.io/apiserver/pkg/util/feature"
 	"k8s.io/klog/v2"
 	"k8s.io/kubernetes/pkg/controller"
+	"k8s.io/kubernetes/pkg/features"
 )
 
 const (
@@ -294,11 +296,17 @@ func addCompletionIndexEnvVariable(container *v1.Container) {
 			return
 		}
 	}
+	var fieldPath string
+	if feature.DefaultFeatureGate.Enabled(features.PodIndexLabel) {
+		fieldPath = fmt.Sprintf("metadata.labels['%s']", batch.JobCompletionIndexAnnotation)
+	} else {
+		fieldPath = fmt.Sprintf("metadata.annotations['%s']", batch.JobCompletionIndexAnnotation)
+	}
 	container.Env = append(container.Env, v1.EnvVar{
 		Name: completionIndexEnvName,
 		ValueFrom: &v1.EnvVarSource{
 			FieldRef: &v1.ObjectFieldSelector{
-				FieldPath: fmt.Sprintf("metadata.labels['%s']", batch.JobCompletionIndexAnnotation),
+				FieldPath: fieldPath,
 			},
 		},
 	})
diff --git a/pkg/controller/job/job_controller_test.go b/pkg/controller/job/job_controller_test.go
index c9c147c6e41..a830a4c7e7b 100644
--- a/pkg/controller/job/job_controller_test.go
+++ b/pkg/controller/job/job_controller_test.go
@@ -973,7 +973,7 @@ func checkIndexedJobPods(t *testing.T, control *controller.FakePodControl, wantI
 	t.Helper()
 	gotIndexes := sets.New[int]()
 	for _, p := range control.Templates {
-		checkJobCompletionEnvVariable(t, &p.Spec)
+		checkJobCompletionEnvVariable(t, &p.Spec, podIndexLabelDisabled)
 		if !podIndexLabelDisabled {
 			checkJobCompletionLabel(t, &p)
 		}
@@ -4418,14 +4418,20 @@ func checkJobCompletionLabel(t *testing.T, p *v1.PodTemplateSpec) {
 	}
 }
 
-func checkJobCompletionEnvVariable(t *testing.T, spec *v1.PodSpec) {
+func checkJobCompletionEnvVariable(t *testing.T, spec *v1.PodSpec, podIndexLabelDisabled bool) {
 	t.Helper()
+	var fieldPath string
+	if podIndexLabelDisabled {
+		fieldPath = fmt.Sprintf("metadata.annotations['%s']", batch.JobCompletionIndexAnnotation)
+	} else {
+		fieldPath = fmt.Sprintf("metadata.labels['%s']", batch.JobCompletionIndexAnnotation)
+	}
 	want := []v1.EnvVar{
 		{
 			Name: "JOB_COMPLETION_INDEX",
 			ValueFrom: &v1.EnvVarSource{
 				FieldRef: &v1.ObjectFieldSelector{
-					FieldPath: fmt.Sprintf("metadata.annotations['%s']", batch.JobCompletionIndexAnnotation),
+					FieldPath: fieldPath,
 				},
 			},
 		},