diff --git a/pkg/admission/admission_controller.go b/pkg/admission/admission_controller.go index 2c75918c0..5c3b5f405 100644 --- a/pkg/admission/admission_controller.go +++ b/pkg/admission/admission_controller.go @@ -425,59 +425,41 @@ func (c *AdmissionController) updateApplicationInfo(namespace string, pod *v1.Po newLabels, newAnnotations := getNewApplicationInfo(pod, namespace, c.conf.GetGenerateUniqueAppIds(), c.conf.GetDefaultQueueName()) - patch = updateLabelPatch(pod, newLabels, patch) - patch = updateAnnotationPatch(pod, newAnnotations, patch) + patch = updatePatch(pod, newLabels, constants.LabelPatchPath, patch) + patch = updatePatch(pod, newAnnotations, constants.AnnotationPatchPath, patch) return patch } -func updateLabelPatch(pod *v1.Pod, labels map[string]string, patch []common.PatchOperation) []common.PatchOperation { - // if found an existing patch on labels, add new labels to it +func updatePatch(pod *v1.Pod, newValues map[string]string, patchPath string, patch []common.PatchOperation) []common.PatchOperation { + // if found an existing patch, add new values to it for _, p := range patch { - if p.Op == constants.AddPatchOp && p.Path == constants.LabelPatchPath { - if existingLabels, ok := p.Value.(map[string]string); ok { - for k, v := range labels { - existingLabels[k] = v + if p.Op == constants.AddPatchOp && p.Path == patchPath { + if existingPatchValues, ok := p.Value.(map[string]string); ok { + for k, v := range newValues { + existingPatchValues[k] = v } return patch } } } - // Add a new label patch - if len(labels) != 0 { - // combine new labels with existing labels in pod to create first patch - existingLabels := pod.Labels - labels = utils.MergeMaps(existingLabels, labels) - patch = append(patch, common.PatchOperation{ - Op: constants.AddPatchOp, - Path: constants.LabelPatchPath, - Value: labels, - }) - } - return patch -} -func updateAnnotationPatch(pod *v1.Pod, annotations map[string]string, patch []common.PatchOperation) []common.PatchOperation { - // if found an existing patch on annotations, add new annotations to it - for _, p := range patch { - if p.Op == constants.AddPatchOp && p.Path == constants.AnnotationPatchPath { - if existingAnnotations, ok := p.Value.(map[string]string); ok { - for k, v := range annotations { - existingAnnotations[k] = v - } - return patch - } + // Add a new patch + if len(newValues) != 0 { + // combine new values with existing values in pod to create first patch + var existingPodValues map[string]string + if patchPath == constants.LabelPatchPath { + // newly add labels patch should include existing labels in pod + existingPodValues = pod.Labels + } else if patchPath == constants.AnnotationPatchPath { + // newly add annotations patch should include existing annotations in pod + existingPodValues = pod.Annotations } - } - // Add a new annotation patch - if len(annotations) != 0 { - // combine new annotations with existing annotations in pod to create first patch - existingAnnotations := pod.Annotations - annotations = utils.MergeMaps(existingAnnotations, annotations) + patchValue := utils.MergeMaps(existingPodValues, newValues) patch = append(patch, common.PatchOperation{ Op: constants.AddPatchOp, - Path: constants.AnnotationPatchPath, - Value: annotations, + Path: patchPath, + Value: patchValue, }) } return patch diff --git a/pkg/admission/admission_controller_test.go b/pkg/admission/admission_controller_test.go index 146729503..d1bf9a68e 100644 --- a/pkg/admission/admission_controller_test.go +++ b/pkg/admission/admission_controller_test.go @@ -231,100 +231,114 @@ func TestUpdateApplicationInfo(t *testing.T) { assert.Equal(t, strings.HasPrefix(annotationsInPatch["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) } -func TestUpdateLabelPatch(t *testing.T) { +func TestUpdatePatch(t *testing.T) { + // dummy pod for following test cases pod := &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ Labels: map[string]string{ "label_key_in_pod": "label_value_in_pod", }, - }, - } - - // Verify if no patch on labels in patch list, - // create a new patch with pods's existing labels and put new labels into the patch - patch := make([]common.PatchOperation, 0) - newLabels := map[string]string{ - "new_key_1": "new_value_1", - "new_key_2": "new_value_2", - } - patch = updateLabelPatch(pod, newLabels, patch) - result := fetchPatchValues(constants.AddPatchOp, constants.LabelPatchPath, patch) - - assert.Equal(t, len(result), 3) - assert.Equal(t, result["label_key_in_pod"], "label_value_in_pod") - assert.Equal(t, result["new_key_1"], "new_value_1") - assert.Equal(t, result["new_key_2"], "new_value_2") - - // Verify if have patch on labels in patch list, - // won't crete new patch and will add new labels into the patch - patch = make([]common.PatchOperation, 0) - patch = append(patch, common.PatchOperation{ - Op: constants.AddPatchOp, - Path: constants.LabelPatchPath, - Value: map[string]string{ - "label_key_in_patch": "label_value_in_patch", - }, - }) - newLabels = map[string]string{ - "new_key_1": "new_value_1", - "new_key_2": "new_value_2", - } - - patch = updateLabelPatch(pod, newLabels, patch) - result = fetchPatchValues(constants.AddPatchOp, constants.LabelPatchPath, patch) - - assert.Equal(t, len(result), 3) - assert.Equal(t, result["label_key_in_patch"], "label_value_in_patch") - assert.Equal(t, result["new_key_1"], "new_value_1") - assert.Equal(t, result["new_key_2"], "new_value_2") -} - -func TestUpdateAnnotationPatch(t *testing.T) { - pod := &v1.Pod{ - ObjectMeta: metav1.ObjectMeta{ Annotations: map[string]string{ "annotation_key_in_pod": "annotation_value_in_pod", }, }, } - // Verify if no patch on annotations in patch list, - // create a new patch with pods's existing annotations and put new annotation into the patch - patch := make([]common.PatchOperation, 0) - newAnnotations := map[string]string{ - "new_key_1": "new_value_1", - "new_key_2": "new_value_2", - } - patch = updateAnnotationPatch(pod, newAnnotations, patch) - result := fetchPatchValues(constants.AddPatchOp, constants.AnnotationPatchPath, patch) - - assert.Equal(t, len(result), 3) - assert.Equal(t, result["annotation_key_in_pod"], "annotation_value_in_pod") - assert.Equal(t, result["new_key_1"], "new_value_1") - assert.Equal(t, result["new_key_2"], "new_value_2") - - // Verify if have patch on annotation in patch list, - // won't crete new patch and will add new annotation into the patch - patch = make([]common.PatchOperation, 0) - patch = append(patch, common.PatchOperation{ - Op: constants.AddPatchOp, - Path: constants.AnnotationPatchPath, - Value: map[string]string{ - "annotation_key_in_patch": "annotation_value_in_patch", + testCases := []struct { + description string + pod *v1.Pod + patch []common.PatchOperation + patchPath string + newValues map[string]string + expectedPatchValues map[string]string + expectedPatchSize int + }{ + { + description: "Verify if no patch on labels in patch list, create new patch with pods's existing labels and put new labels into the patch", + pod: pod, + patch: make([]common.PatchOperation, 0), + patchPath: constants.LabelPatchPath, + newValues: map[string]string{ + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchValues: map[string]string{ + "label_key_in_pod": "label_value_in_pod", + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchSize: 1, + }, { + description: "Verify if no patch on annotations in patch list, create new patch with pods's existing annotation and put new annotation into the patch", + pod: pod, + patch: make([]common.PatchOperation, 0), + patchPath: constants.AnnotationPatchPath, + newValues: map[string]string{ + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchValues: map[string]string{ + "annotation_key_in_pod": "annotation_value_in_pod", + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchSize: 1, + }, { + description: "Verify if have patch on labels in patch list, won't crete new patch and will add new labels into the existing patch", + pod: pod, + patch: []common.PatchOperation{ + { + Op: constants.AddPatchOp, + Path: constants.LabelPatchPath, + Value: map[string]string{ + "label_key_in_patch": "label_value_in_patch", + }, + }, + }, + patchPath: constants.LabelPatchPath, + newValues: map[string]string{ + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchValues: map[string]string{ + "label_key_in_patch": "label_value_in_patch", + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchSize: 1, + }, { + description: "Verify if have patch on annotations in patch list, won't crete new patch and will add new annotations into the existing patch", + pod: pod, + patch: []common.PatchOperation{ + { + Op: constants.AddPatchOp, + Path: constants.AnnotationPatchPath, + Value: map[string]string{ + "annotation_key_in_patch": "annotation_value_in_patch", + }, + }, + }, + patchPath: constants.AnnotationPatchPath, + newValues: map[string]string{ + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchValues: map[string]string{ + "annotation_key_in_patch": "annotation_value_in_patch", + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchSize: 1, }, - }) - newAnnotations = map[string]string{ - "new_key_1": "new_value_1", - "new_key_2": "new_value_2", } - - patch = updateAnnotationPatch(pod, newAnnotations, patch) - result = fetchPatchValues(constants.AddPatchOp, constants.AnnotationPatchPath, patch) - - assert.Equal(t, len(result), 3) - assert.Equal(t, result["annotation_key_in_patch"], "annotation_value_in_patch") - assert.Equal(t, result["new_key_1"], "new_value_1") - assert.Equal(t, result["new_key_2"], "new_value_2") + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + patch := updatePatch(tc.pod, tc.newValues, tc.patchPath, tc.patch) + assert.DeepEqual(t, len(patch), tc.expectedPatchSize) + assert.DeepEqual(t, patch[0].Path, tc.patchPath) + assert.DeepEqual(t, patch[0].Value, tc.expectedPatchValues) + }) + } } func fetchPatchValues(op string, path string, patch []common.PatchOperation) map[string]string { diff --git a/pkg/admission/util.go b/pkg/admission/util.go index 4c21ae787..bf5a1d42f 100644 --- a/pkg/admission/util.go +++ b/pkg/admission/util.go @@ -64,7 +64,7 @@ func getNewApplicationInfo(pod *v1.Pod, namespace string, generateUniqueAppIds b } // we're looking forward to deprecate the labels and move everything to annotations - // but for backforward compatiblity, we still add to labels + // but for backforward compatibility, we still add to labels newLabels = updateLabelIfNotPresentInPod(pod, newLabels, constants.LabelApplicationID, appID) newAnnotations = updateAnnotationIfNotPresentInPod(pod, newAnnotations, constants.AnnotationApplicationID, appID)