diff --git a/pkg/kube/pod.go b/pkg/kube/pod.go index 3c772b77af..339189110a 100644 --- 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 index 3b84a74d75..d8527d4ac4 100644 --- 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) + } +}