From 7d3b7d9935c82f84e63c04283d52902ce0df8ce6 Mon Sep 17 00:00:00 2001
From: Aldo Culquicondor <1299064+alculquicondor@users.noreply.github.com>
Date: Thu, 12 Aug 2021 10:58:03 -0400
Subject: [PATCH] Fix ReplicaType key for filtering pods and services (#152)

---
 pkg/controller.v1/common/pod.go          |  2 +-
 pkg/controller.v1/common/pod_test.go     | 45 ++++++++++++++++++++++++
 pkg/controller.v1/common/service.go      |  2 +-
 pkg/controller.v1/common/service_test.go | 44 +++++++++++++++++++++++
 4 files changed, 91 insertions(+), 2 deletions(-)

diff --git a/pkg/controller.v1/common/pod.go b/pkg/controller.v1/common/pod.go
index e2d00fdc..ffbce85b 100644
--- a/pkg/controller.v1/common/pod.go
+++ b/pkg/controller.v1/common/pod.go
@@ -258,7 +258,7 @@ func (jc *JobController) FilterPodsForReplicaType(pods []*v1.Pod, replicaType ap
 	var result []*v1.Pod
 
 	selector := labels.SelectorFromValidatedSet(labels.Set{
-		apiv1.ReplicaIndexLabel: string(replicaType),
+		apiv1.ReplicaTypeLabel: string(replicaType),
 	})
 
 	// TODO(#149): Remove deprecated selector.
diff --git a/pkg/controller.v1/common/pod_test.go b/pkg/controller.v1/common/pod_test.go
index 8f390f9c..ccd5fae6 100644
--- a/pkg/controller.v1/common/pod_test.go
+++ b/pkg/controller.v1/common/pod_test.go
@@ -146,3 +146,48 @@ func TestCalculatePodSliceSize(t *testing.T) {
 		assert.Equal(t, tc.expectedSize, result)
 	}
 }
+
+func TestFilterPodsForReplicaType(t *testing.T) {
+	pods := []*v1.Pod{
+		{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:   "a",
+				Labels: map[string]string{apiv1.ReplicaTypeLabel: "foo"},
+			},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:   "b",
+				Labels: map[string]string{apiv1.ReplicaTypeLabel: "bar"},
+			},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:   "c",
+				Labels: map[string]string{apiv1.ReplicaTypeLabelDeprecated: "foo"},
+			},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:   "d",
+				Labels: map[string]string{apiv1.ReplicaTypeLabelDeprecated: "bar"},
+			},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "e",
+				Labels: map[string]string{
+					apiv1.ReplicaTypeLabel:           "foo",
+					apiv1.ReplicaTypeLabelDeprecated: "bar",
+				},
+			},
+		},
+	}
+	c := &JobController{}
+	got, err := c.FilterPodsForReplicaType(pods, "foo")
+	if err != nil {
+		t.Fatalf("FilterPodsForReplicaType returned error: %v", err)
+	}
+	want := []*v1.Pod{pods[0], pods[2], pods[4]}
+	assert.Equal(t, want, got)
+}
diff --git a/pkg/controller.v1/common/service.go b/pkg/controller.v1/common/service.go
index f51de0e4..f85e0f4b 100644
--- a/pkg/controller.v1/common/service.go
+++ b/pkg/controller.v1/common/service.go
@@ -142,7 +142,7 @@ func (jc *JobController) FilterServicesForReplicaType(services []*v1.Service, re
 	var result []*v1.Service
 
 	selector := labels.SelectorFromValidatedSet(labels.Set{
-		apiv1.ReplicaIndexLabel: string(replicaType),
+		apiv1.ReplicaTypeLabel: string(replicaType),
 	})
 
 	// TODO(#149): Remove deprecated selector.
diff --git a/pkg/controller.v1/common/service_test.go b/pkg/controller.v1/common/service_test.go
index 68392f26..2592be99 100644
--- a/pkg/controller.v1/common/service_test.go
+++ b/pkg/controller.v1/common/service_test.go
@@ -208,3 +208,47 @@ func TestJobController_CreateNewService(t *testing.T) {
 		})
 	}
 }
+func TestFilterServicesForReplicaType(t *testing.T) {
+	services := []*v1.Service{
+		{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:   "a",
+				Labels: map[string]string{apiv1.ReplicaTypeLabel: "foo"},
+			},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:   "b",
+				Labels: map[string]string{apiv1.ReplicaTypeLabel: "bar"},
+			},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:   "c",
+				Labels: map[string]string{apiv1.ReplicaTypeLabelDeprecated: "foo"},
+			},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:   "d",
+				Labels: map[string]string{apiv1.ReplicaTypeLabelDeprecated: "bar"},
+			},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "e",
+				Labels: map[string]string{
+					apiv1.ReplicaTypeLabel:           "foo",
+					apiv1.ReplicaTypeLabelDeprecated: "bar",
+				},
+			},
+		},
+	}
+	c := &JobController{}
+	got, err := c.FilterServicesForReplicaType(services, "foo")
+	if err != nil {
+		t.Fatalf("FilterPodsForReplicaType returned error: %v", err)
+	}
+	want := []*v1.Service{services[0], services[2], services[4]}
+	assert.Equal(t, want, got)
+}