Add AddLabels and AddAnnotations method to PodOptions type (#3073)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
kanisterio · Aug 28, 2024 · ddbe4d1 · ddbe4d1
1 parent 3183273
commit ddbe4d1
Show file tree

Hide file tree

Showing 2 changed files with 190 additions and 0 deletions.
diff --git a/pkg/kube/pod.go b/pkg/kube/pod.go
@@ -84,6 +84,34 @@ type PodOptions struct {
 	Lifecycle                *corev1.Lifecycle
 }
 
+func (po *PodOptions) AddLabels(labels map[string]string) {
+	if po == nil {
+		return
+	}
+
+	if po.Labels == nil {
+		po.Labels = make(map[string]string)
+	}
+
+	for k, v := range labels {
+		po.Labels[k] = v
+	}
+}
+
+func (po *PodOptions) AddAnnotations(annotations map[string]string) {
+	if po == nil {
+		return
+	}
+
+	if po.Annotations == nil {
+		po.Annotations = make(map[string]string)
+	}
+
+	for k, v := range annotations {
+		po.Annotations[k] = v
+	}
+}
+
 func GetPodObjectFromPodOptions(ctx context.Context, cli kubernetes.Interface, opts *PodOptions) (*corev1.Pod, error) {
 	// If Namespace is not specified, use the controller Namespace.
 	cns, err := GetControllerNamespace()

diff --git a/pkg/kube/pod_test.go b/pkg/kube/pod_test.go
@@ -1078,3 +1078,165 @@ func (s *PodControllerTestSuite) TestContainerNameFromPodOptsOrDefault(c *C) {
 	name = ContainerNameFromPodOptsOrDefault(nil)
 	c.Assert(name, Equals, DefaultContainerName)
 }
+
+func (s *PodSuite) TestAddLabels(c *C) {
+	for _, tc := range []struct {
+		podOptions         *PodOptions
+		labels             map[string]string
+		expectedPodOptions *PodOptions
+	}{
+		{
+			podOptions: &PodOptions{},
+			labels: map[string]string{
+				"keyOne": "valOne",
+			},
+			expectedPodOptions: &PodOptions{
+				Labels: map[string]string{
+					"keyOne": "valOne",
+				},
+			},
+		},
+		{
+			podOptions: nil,
+			labels: map[string]string{
+				"keyOne": "valOne",
+			},
+			expectedPodOptions: nil,
+		},
+		{
+			podOptions: &PodOptions{
+				Labels: map[string]string{
+					"key": "val",
+				},
+			},
+			labels: map[string]string{
+				"keyOne": "valOne",
+			},
+			expectedPodOptions: &PodOptions{
+				Labels: map[string]string{
+					"key":    "val",
+					"keyOne": "valOne",
+				},
+			},
+		},
+		{
+			podOptions: &PodOptions{
+				Labels: map[string]string{
+					"key":     "val",
+					"keyZero": "valZero",
+				},
+			},
+			labels: map[string]string{
+				"keyOne": "valOne",
+				"keyTwo": "valTwo",
+			},
+			expectedPodOptions: &PodOptions{
+				Labels: map[string]string{
+					"key":     "val",
+					"keyZero": "valZero",
+					"keyOne":  "valOne",
+					"keyTwo":  "valTwo",
+				},
+			},
+		},
+		{
+			podOptions: &PodOptions{
+				Labels: map[string]string{
+					"key":     "val",
+					"keyZero": "valZero",
+				},
+			},
+			labels: nil,
+			expectedPodOptions: &PodOptions{
+				Labels: map[string]string{
+					"key":     "val",
+					"keyZero": "valZero",
+				},
+			},
+		},
+	} {
+		tc.podOptions.AddLabels(tc.labels)
+		c.Assert(tc.podOptions, DeepEquals, tc.expectedPodOptions)
+	}
+}
+
+func (s *PodSuite) TestAddAnnotations(c *C) {
+	for _, tc := range []struct {
+		podOptions         *PodOptions
+		annotations        map[string]string
+		expectedPodOptions *PodOptions
+	}{
+		{
+			podOptions: &PodOptions{},
+			annotations: map[string]string{
+				"keyOne": "valOne",
+			},
+			expectedPodOptions: &PodOptions{
+				Annotations: map[string]string{
+					"keyOne": "valOne",
+				},
+			},
+		},
+		{
+			podOptions: nil,
+			annotations: map[string]string{
+				"keyOne": "valOne",
+			},
+			expectedPodOptions: nil,
+		},
+		{
+			podOptions: &PodOptions{
+				Annotations: map[string]string{
+					"key": "val",
+				},
+			},
+			annotations: map[string]string{
+				"keyOne": "valOne",
+			},
+			expectedPodOptions: &PodOptions{
+				Annotations: map[string]string{
+					"key":    "val",
+					"keyOne": "valOne",
+				},
+			},
+		},
+		{
+			podOptions: &PodOptions{
+				Annotations: map[string]string{
+					"key":     "val",
+					"keyZero": "valZero",
+				},
+			},
+			annotations: map[string]string{
+				"keyOne": "valOne",
+				"keyTwo": "valTwo",
+			},
+			expectedPodOptions: &PodOptions{
+				Annotations: map[string]string{
+					"key":     "val",
+					"keyZero": "valZero",
+					"keyOne":  "valOne",
+					"keyTwo":  "valTwo",
+				},
+			},
+		},
+		{
+			podOptions: &PodOptions{
+				Annotations: map[string]string{
+					"key":     "val",
+					"keyZero": "valZero",
+				},
+			},
+			annotations: nil,
+			expectedPodOptions: &PodOptions{
+				Annotations: map[string]string{
+					"key":     "val",
+					"keyZero": "valZero",
+				},
+			},
+		},
+	} {
+		tc.podOptions.AddAnnotations(tc.annotations)
+		c.Assert(tc.podOptions, DeepEquals, tc.expectedPodOptions)
+	}
+}