From be5421f00b2d9eac968259a0aea41c58da537669 Mon Sep 17 00:00:00 2001
From: f10atin9 <f10atin9@kubesphere.io>
Date: Mon, 27 Sep 2021 16:42:19 +0800
Subject: [PATCH] update capability_controller.go, make sure that annotations
 is generated correctly.StorageClass without csiDriver will no longer generate
 false annotations.

Signed-off-by: f10atin9 <f10atin9@kubesphere.io>
---
 .../capability/capability_controller.go       | 27 ++++++++++++-------
 .../capability/capability_controller_test.go  | 19 +++++++++++++
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/pkg/controller/storage/capability/capability_controller.go b/pkg/controller/storage/capability/capability_controller.go
index 27cc7b2ef..4c5b209ab 100644
--- a/pkg/controller/storage/capability/capability_controller.go
+++ b/pkg/controller/storage/capability/capability_controller.go
@@ -208,13 +208,17 @@ func (c *StorageCapabilityController) syncHandler(key string) error {
 
 	// Annotate storageClass
 	storageClassUpdated := storageClass.DeepCopy()
-	err = c.addStorageClassSnapshotAnnotation(storageClassUpdated, isCSIStorage)
-	if err != nil {
-		return err
-	}
-	err = c.addCloneVolumeAnnotation(storageClassUpdated, isCSIStorage)
-	if err != nil {
-		return err
+	if !isCSIStorage {
+		c.removeAnnotations(storageClassUpdated)
+	} else {
+		err = c.updateStorageClassSnapshotAnnotation(storageClassUpdated, isCSIStorage)
+		if err != nil {
+			return err
+		}
+		err = c.updateCloneVolumeAnnotation(storageClassUpdated, isCSIStorage)
+		if err != nil {
+			return err
+		}
 	}
 	_, err = c.storageClassClient.Update(context.Background(), storageClassUpdated, metav1.UpdateOptions{})
 	if err != nil {
@@ -234,7 +238,7 @@ func (c *StorageCapabilityController) hasCSIDriver(storageClass *storagev1.Stora
 	return false
 }
 
-func (c *StorageCapabilityController) addStorageClassSnapshotAnnotation(storageClass *storagev1.StorageClass, snapshotAllow bool) error {
+func (c *StorageCapabilityController) updateStorageClassSnapshotAnnotation(storageClass *storagev1.StorageClass, snapshotAllow bool) error {
 	if storageClass.Annotations == nil {
 		storageClass.Annotations = make(map[string]string)
 	}
@@ -246,7 +250,7 @@ func (c *StorageCapabilityController) addStorageClassSnapshotAnnotation(storageC
 	return nil
 }
 
-func (c *StorageCapabilityController) addCloneVolumeAnnotation(storageClass *storagev1.StorageClass, cloneAllow bool) error {
+func (c *StorageCapabilityController) updateCloneVolumeAnnotation(storageClass *storagev1.StorageClass, cloneAllow bool) error {
 	if storageClass.Annotations == nil {
 		storageClass.Annotations = make(map[string]string)
 	}
@@ -256,3 +260,8 @@ func (c *StorageCapabilityController) addCloneVolumeAnnotation(storageClass *sto
 	}
 	return nil
 }
+
+func (c *StorageCapabilityController) removeAnnotations(storageClass *storagev1.StorageClass) {
+	delete(storageClass.Annotations, annotationAllowClone)
+	delete(storageClass.Annotations, annotationAllowSnapshot)
+}
diff --git a/pkg/controller/storage/capability/capability_controller_test.go b/pkg/controller/storage/capability/capability_controller_test.go
index 22486e5e5..7890d5d2f 100644
--- a/pkg/controller/storage/capability/capability_controller_test.go
+++ b/pkg/controller/storage/capability/capability_controller_test.go
@@ -270,3 +270,22 @@ func TestStorageClassHadOneAnnotation(t *testing.T) {
 	// Run test
 	fixture.run(getKey(storageClass, t))
 }
+
+func TestStorageClassHadNoCSIDriver(t *testing.T) {
+	fixture := newFixture(t, true)
+	storageClass := newStorageClass("csi-example", "csi.example.com")
+	storageClass.Annotations = map[string]string{}
+	storageClassUpdate := storageClass.DeepCopy()
+	storageClass.Annotations = map[string]string{annotationAllowSnapshot: "false"}
+	storageClass.Annotations = map[string]string{annotationAllowClone: "false"}
+
+	// Object exist
+	fixture.storageObjects = append(fixture.storageObjects, storageClass)
+	fixture.storageClassLister = append(fixture.storageClassLister, storageClass)
+
+	// Action expected
+	fixture.expectUpdateStorageClassAction(storageClassUpdate)
+
+	// Run test
+	fixture.run(getKey(storageClass, t))
+}