From cb91d2efe26f3325fd34c57738641956d2271394 Mon Sep 17 00:00:00 2001
From: Jingfang Liu <jingfangliu@google.com>
Date: Wed, 13 Jan 2021 10:33:42 -0800
Subject: [PATCH] Add generic function for updating the annotations only
 (#1351)

---
 commands/migratecmd.go    |  2 +-
 pkg/client/client.go      | 32 ++++++++++++--------
 pkg/client/client_test.go | 62 +++++++++++++++++++++++++++++++++++++--
 3 files changed, 80 insertions(+), 16 deletions(-)

diff --git a/commands/migratecmd.go b/commands/migratecmd.go
index d408841723..37b8f368d0 100644
--- a/commands/migratecmd.go
+++ b/commands/migratecmd.go
@@ -354,7 +354,7 @@ func updateOwningInventoryAnnotation(f cmdutil.Factory, objMetas []object.ObjMet
 			}
 			return err
 		}
-		changed, err := client.UpdateAnnotation(obj, old, new)
+		changed, err := client.ReplaceOwningInventoryID(obj, old, new)
 		if err != nil {
 			return err
 		}
diff --git a/pkg/client/client.go b/pkg/client/client.go
index 662c79e1b2..24b4dbc578 100644
--- a/pkg/client/client.go
+++ b/pkg/client/client.go
@@ -61,10 +61,10 @@ func (uc *client) resourceInterface(meta object.ObjMetadata) (dynamic.ResourceIn
 	return namespacedClient, nil
 }
 
-// UpdateAnnotation updates the object owning inventory annotation
+// ReplaceOwningInventoryID updates the object owning inventory annotation
 // to the new ID when the owning inventory annotation is either empty or the old ID.
-// It returns if the annotation is updated.
-func UpdateAnnotation(obj *unstructured.Unstructured, oldID, newID string) (bool, error) {
+// It returns true if the annotation is updated.
+func ReplaceOwningInventoryID(obj *unstructured.Unstructured, oldID, newID string) (bool, error) {
 	key := "config.k8s.io/owning-inventory"
 	annotations := obj.GetAnnotations()
 	if annotations == nil {
@@ -73,19 +73,25 @@ func UpdateAnnotation(obj *unstructured.Unstructured, oldID, newID string) (bool
 	val, found := annotations[key]
 	if !found || val == oldID {
 		annotations[key] = newID
+		return true, UpdateAnnotation(obj, annotations)
+	}
+	return false, nil
+}
+
+// UpdateAnnotation updates .metadata.annotations field of obj to use the passed in annotations
+// as well as updates the last-applied-configuration annotations.
+func UpdateAnnotation(obj *unstructured.Unstructured, annotations map[string]string) error {
+	u := getOriginalObj(obj)
+	if u != nil {
+		u.SetAnnotations(annotations)
 		// Since the annotation is updated, we also need to update the
 		// last applied configuration annotation.
-		u := getOriginalObj(obj)
-		if u != nil {
-			u.SetAnnotations(annotations)
-			err := util.CreateOrUpdateAnnotation(false, u, scheme.DefaultJSONEncoder())
-			obj.SetAnnotations(u.GetAnnotations())
-			return true, err
-		}
-		obj.SetAnnotations(annotations)
-		return true, nil
+		err := util.CreateOrUpdateAnnotation(true, u, scheme.DefaultJSONEncoder())
+		obj.SetAnnotations(u.GetAnnotations())
+		return err
 	}
-	return false, nil
+	obj.SetAnnotations(annotations)
+	return nil
 }
 
 func getOriginalObj(obj *unstructured.Unstructured) *unstructured.Unstructured {
diff --git a/pkg/client/client_test.go b/pkg/client/client_test.go
index efa2b9e936..0b9843f4a7 100644
--- a/pkg/client/client_test.go
+++ b/pkg/client/client_test.go
@@ -4,12 +4,15 @@
 package client
 
 import (
+	"reflect"
 	"testing"
 
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/kubectl/pkg/scheme"
+	"k8s.io/kubectl/pkg/util"
 )
 
-func TestUpdateAnnotation(t *testing.T) {
+func TestReplaceOwningInventoryID(t *testing.T) {
 	testcases := []struct {
 		name         string
 		annotations  map[string]string
@@ -55,7 +58,7 @@ func TestUpdateAnnotation(t *testing.T) {
 	}
 	for _, tc := range testcases {
 		deployment.SetAnnotations(tc.annotations)
-		updated, err := UpdateAnnotation(deployment, tc.oldID, tc.newID)
+		updated, err := ReplaceOwningInventoryID(deployment, tc.oldID, tc.newID)
 		if err != nil {
 			t.Errorf("unexpected error %v", err)
 		}
@@ -71,3 +74,58 @@ func TestUpdateAnnotation(t *testing.T) {
 		}
 	}
 }
+
+func TestUpdateAnnotation(t *testing.T) {
+	testcases := []struct {
+		name        string
+		annotations map[string]string
+	}{
+		{
+			name: "add new annotations",
+			annotations: map[string]string{
+				"config.k8s.io/owning-inventory": "new",
+				"random-key":                     "value",
+			},
+		},
+		{
+			name: "remove existing annotations",
+			annotations: map[string]string{
+				"random-key": "value",
+			},
+		},
+	}
+
+	for _, tc := range testcases {
+		u := &unstructured.Unstructured{
+			Object: map[string]interface{}{
+				"apiVersion": "apps/v1",
+				"kind":       "Deployment",
+				"metadata": map[string]interface{}{
+					"name":      "deployment",
+					"namespace": "test",
+					"annotations": map[string]interface{}{
+						"config.k8s.io/owning-inventory": "old",
+					},
+				},
+			},
+		}
+		uCloned := u.DeepCopy()
+		uCloned.SetAnnotations(tc.annotations)
+		err := util.CreateOrUpdateAnnotation(true, uCloned, scheme.DefaultJSONEncoder())
+		if err != nil {
+			t.Errorf("unexpected error %v", err)
+		}
+		err = util.CreateOrUpdateAnnotation(true, u, scheme.DefaultJSONEncoder())
+		if err != nil {
+			t.Errorf("unexpected error %v", err)
+		}
+		err = UpdateAnnotation(u, tc.annotations)
+		if err != nil {
+			t.Errorf("unexpected error %v", err)
+		}
+
+		if !reflect.DeepEqual(u, uCloned) {
+			t.Errorf("%s failed: expected %v, but got %v", tc.name, uCloned, u)
+		}
+	}
+}