chore(api): simplify annotation clearing

Signed-off-by: Hidde Beydals <hidde@hhh.computer>
akuity · Mar 15, 2024 · e0629ca · e0629ca
1 parent bdf2aad
commit e0629ca
Show file tree

Hide file tree

Showing 9 changed files with 167 additions and 317 deletions.
diff --git a/api/v1alpha1/helpers.go b/api/v1alpha1/helpers.go
@@ -8,6 +8,26 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
+func ClearAnnotations(ctx context.Context, c client.Client, obj client.Object, keys ...string) error {
+	if len(keys) == 0 {
+		return nil
+	}
+
+	patchBytes := []byte(`{"metadata":{"annotations":{`)
+	for i, key := range keys {
+		if i > 0 {
+			patchBytes = append(patchBytes, ',')
+		}
+		patchBytes = append(patchBytes, fmt.Sprintf(`"%s":null`, key)...)
+	}
+	patchBytes = append(patchBytes, "}}}"...)
+	patch := client.RawPatch(types.MergePatchType, patchBytes)
+	if err := c.Patch(ctx, obj, patch); err != nil {
+		return fmt.Errorf("patch annotation: %w", err)
+	}
+	return nil
+}
+
 func patchAnnotation(ctx context.Context, c client.Client, obj client.Object, key, value string) error {
 	patchBytes := []byte(
 		fmt.Sprintf(
@@ -22,17 +42,3 @@ func patchAnnotation(ctx context.Context, c client.Client, obj client.Object, ke
 	}
 	return nil
 }
-
-func clearObjectAnnotation(
-	ctx context.Context,
-	c client.Client,
-	obj client.Object,
-	annotationKey string,
-) error {
-	patchBytes := []byte(fmt.Sprintf(`{"metadata":{"annotations":{"%s":null}}}`, annotationKey))
-	patch := client.RawPatch(types.MergePatchType, patchBytes)
-	if err := c.Patch(ctx, obj, patch); err != nil {
-		return fmt.Errorf("patch annotation: %w", err)
-	}
-	return nil
-}
diff --git a/api/v1alpha1/helpers_test.go b/api/v1alpha1/helpers_test.go
@@ -11,6 +11,128 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 )
 
+func TestClearAnnotations(t *testing.T) {
+	scheme := k8sruntime.NewScheme()
+	require.NoError(t, SchemeBuilder.AddToScheme(scheme))
+
+	t.Parallel()
+	newFakeClient := func(obj ...client.Object) client.Client {
+		return fake.NewClientBuilder().
+			WithScheme(scheme).
+			WithObjects(obj...).
+			Build()
+	}
+
+	testCases := []struct{
+		name string
+		client client.Client
+		obj client.Object
+		keys []string
+		assertions func(*testing.T, client.Object, error)
+	}{
+		{
+			name: "no keys",
+			client: newFakeClient(),
+			obj: nil,
+			keys: nil,
+			assertions: func(t *testing.T, _ client.Object, err error) {
+				require.NoError(t, err)
+			},
+		},
+		{
+			name: "no annotations",
+			client: newFakeClient(&Stage{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: "test",
+					Name: "stage",
+				},
+			}),
+			obj: &Stage{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: "test",
+					Name: "stage",
+				},
+			},
+			keys: []string{"key"},
+			assertions: func(t *testing.T, _ client.Object, err error) {
+				require.NoError(t, err)
+			},
+		},
+		{
+			name: "not found",
+			client: newFakeClient(),
+			obj: &Stage{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: "test",
+					Name: "stage",
+				},
+			},
+			keys: []string{"key"},
+			assertions: func(t *testing.T, _ client.Object, err error) {
+				require.ErrorContains(t, err, "patch annotation")
+			},
+		},
+		{
+			name: "clear one",
+			client: newFakeClient(&Stage{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: "test",
+					Name: "stage",
+					Annotations: map[string]string{
+						"key1": "value1",
+						"key2": "value2",
+					},
+				},
+			}),
+			obj: &Stage{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: "test",
+					Name: "stage",
+				},
+			},
+			keys: []string{"key1"},
+			assertions: func(t *testing.T, obj client.Object, err error) {
+				require.NoError(t, err)
+				require.Contains(t, obj.GetAnnotations(), "key2")
+				require.NotContains(t, obj.GetAnnotations(), "key1")
+			},
+		},
+		{
+			name: "clear two",
+			client: newFakeClient(&Stage{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: "test",
+					Name: "stage",
+					Annotations: map[string]string{
+						"key1": "value1",
+						"key2": "value2",
+					},
+				},
+			}),
+			obj: &Stage{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: "test",
+					Name: "stage",
+				},
+			},
+			keys: []string{"key1", "key2"},
+			assertions: func(t *testing.T, obj client.Object, err error) {
+				require.NoError(t, err)
+				require.NotContains(t, obj.GetAnnotations(), "key2")
+				require.NotContains(t, obj.GetAnnotations(), "key1")
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			err := ClearAnnotations(context.TODO(), tc.client, tc.obj, tc.keys...)
+			tc.assertions(t, tc.obj, err)
+		})
+	}
+}
+
 func Test_patchAnnotation(t *testing.T) {
 	scheme := k8sruntime.NewScheme()
 	require.NoError(t, SchemeBuilder.AddToScheme(scheme))

diff --git a/api/v1alpha1/promotion_helpers.go b/api/v1alpha1/promotion_helpers.go
@@ -53,25 +53,3 @@ func RefreshPromotion(
 	}
 	return promo, nil
 }
-
-// ClearPromotionRefresh is called by the Promotion controller to clear the refresh
-// annotation on the Promotion (if present).
-func ClearPromotionRefresh(
-	ctx context.Context,
-	c client.Client,
-	promo *Promotion,
-) error {
-	if promo.Annotations == nil {
-		return nil
-	}
-	if _, ok := promo.Annotations[AnnotationKeyRefresh]; !ok {
-		return nil
-	}
-	newPromo := Promotion{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      promo.Name,
-			Namespace: promo.Namespace,
-		},
-	}
-	return clearObjectAnnotation(ctx, c, &newPromo, AnnotationKeyRefresh)
-}
diff --git a/api/v1alpha1/stage_helpers.go b/api/v1alpha1/stage_helpers.go
@@ -55,30 +55,6 @@ func RefreshStage(
 	return stage, nil
 }
 
-// ClearStageRefresh is called by the Stage controller to clear the refresh
-// annotation on the Stage (if present). A client (e.g. UI) who requested a
-// Stage refresh, can wait until the annotation is cleared, to understand that
-// the controller successfully reconciled the Stage after the refresh request.
-func ClearStageRefresh(
-	ctx context.Context,
-	c client.Client,
-	stage *Stage,
-) error {
-	if stage.Annotations == nil {
-		return nil
-	}
-	if _, ok := stage.Annotations[AnnotationKeyRefresh]; !ok {
-		return nil
-	}
-	newStage := Stage{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      stage.Name,
-			Namespace: stage.Namespace,
-		},
-	}
-	return clearObjectAnnotation(ctx, c, &newStage, AnnotationKeyRefresh)
-}
-
 // ReverifyStageFreight forces reconfirmation of the verification of the
 // Freight associated with a Stage by setting an AnnotationKeyReverify
 // annotation on the Stage, causing the controller to rerun the verification.
@@ -111,33 +87,6 @@ func ReverifyStageFreight(
 	return patchAnnotation(ctx, c, stage, AnnotationKeyReverify, curFreight.VerificationInfo.ID)
 }
 
-// ClearStageReverify is called by the Stage controller to clear the
-// AnnotationKeyReverify annotation on the Stage (if present). A client (e.g.
-// UI) who requested a reconfirmation of the Stage verification, can wait
-// until the annotation is cleared, to understand that the controller
-// acknowledged the reverification request.
-func ClearStageReverify(
-	ctx context.Context,
-	c client.Client,
-	stage *Stage,
-) error {
-	if stage.Annotations == nil {
-		return nil
-	}
-
-	if _, ok := stage.Annotations[AnnotationKeyReverify]; !ok {
-		return nil
-	}
-
-	newStage := Stage{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      stage.Name,
-			Namespace: stage.Namespace,
-		},
-	}
-	return clearObjectAnnotation(ctx, c, &newStage, AnnotationKeyReverify)
-}
-
 // AbortStageFreightVerification forces aborting the verification of the
 // Freight associated with a Stage by setting an AnnotationKeyAbort
 // annotation on the Stage, causing the controller to abort the verification.
@@ -174,30 +123,3 @@ func AbortStageFreightVerification(
 
 	return patchAnnotation(ctx, c, stage, AnnotationKeyAbort, curFreight.VerificationInfo.ID)
 }
-
-// ClearStageAbort is called by the Stage controller to clear the
-// AnnotationKeyAbort annotation on the Stage (if present). A client (e.g.
-// UI) who requested an abort of the Stage verification, can wait
-// until the annotation is cleared, to understand that the controller
-// acknowledged the abort request.
-func ClearStageAbort(
-	ctx context.Context,
-	c client.Client,
-	stage *Stage,
-) error {
-	if stage.Annotations == nil {
-		return nil
-	}
-
-	if _, ok := stage.Annotations[AnnotationKeyAbort]; !ok {
-		return nil
-	}
-
-	newStage := Stage{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      stage.Name,
-			Namespace: stage.Namespace,
-		},
-	}
-	return clearObjectAnnotation(ctx, c, &newStage, AnnotationKeyAbort)
-}