Skip to content

Commit

Permalink
combine updateLabelPatch() and updateAnnotationPatch() to updatePatch()
Browse files Browse the repository at this point in the history
  • Loading branch information
@chenyulin0719
chenyulin0719 committed Nov 20, 2023
1 parent cf71eef commit f48e089
Show file tree
Hide file tree
Showing 3 changed files with 117 additions and 121 deletions.
60 changes: 21 additions & 39 deletions pkg/admission/admission_controller.go
View file
Expand Up @@ -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
Expand Down
176 changes: 95 additions & 81 deletions pkg/admission/admission_controller_test.go
View file
Expand Up @@ -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 {
Expand Down
2 changes: 1 addition & 1 deletion pkg/admission/util.go
View file
Expand Up @@ -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)

Expand Down

0 comments on commit f48e089

Please sign in to comment.