Adjust UT with new behavior

Signed-off-by: Shelly Kagan <skagan@redhat.com>
kubevirt · May 22, 2023 · 6338aff · 6338aff
1 parent 2a45f24
commit 6338aff
Showing 1 changed file with 9 additions and 10 deletions.
diff --git a/pkg/controller/populators/upload-populator_test.go b/pkg/controller/populators/upload-populator_test.go
@@ -93,8 +93,7 @@ var _ = Describe("Datavolume controller reconcile loop", func() {
 
 	It("should set event if upload pod failed", func() {
 		pvc := newUploadPopulatorPVC("test-pvc")
-		pvc.Annotations = make(map[string]string)
-		pvc.Annotations[AnnPVCPrimeName] = PVCPrimeName(pvc)
+		cc.AddAnnotation(pvc, AnnPVCPrimeName, PVCPrimeName(pvc))
 		uploadPV := uploadPV(pvc)
 
 		volumeUploadSourceCR := newUploadPopulatorCR("", false)
@@ -106,7 +105,7 @@ var _ = Describe("Datavolume controller reconcile loop", func() {
 		Expect(err).ToNot(HaveOccurred())
 		pvcPrime, err := r.getPVCPrime(pvc)
 		Expect(err).ToNot(HaveOccurred())
-		pvcPrime.Annotations[cc.AnnPodPhase] = string(corev1.PodFailed)
+		cc.AddAnnotation(pvcPrime, cc.AnnPodPhase, string(corev1.PodFailed))
 		pvcPrime.Spec.VolumeName = "test-pv"
 		pvcPrime.UID = pvcPrimeUID
 		err = r.client.Update(context.TODO(), pvcPrime)
@@ -126,8 +125,7 @@ var _ = Describe("Datavolume controller reconcile loop", func() {
 
 	It("should rebind PV to target PVC", func() {
 		pvc := newUploadPopulatorPVC("test-pvc")
-		pvc.Annotations = make(map[string]string)
-		pvc.Annotations[AnnPVCPrimeName] = PVCPrimeName(pvc)
+		cc.AddAnnotation(pvc, AnnPVCPrimeName, PVCPrimeName(pvc))
 		uploadPV := uploadPV(pvc)
 
 		volumeUploadSourceCR := newUploadPopulatorCR("", false)
@@ -139,7 +137,8 @@ var _ = Describe("Datavolume controller reconcile loop", func() {
 		Expect(err).ToNot(HaveOccurred())
 		pvcPrime, err := r.getPVCPrime(pvc)
 		Expect(err).ToNot(HaveOccurred())
-		pvcPrime.Annotations[cc.AnnPodPhase] = string(corev1.PodSucceeded)
+
+		cc.AddAnnotation(pvcPrime, cc.AnnPodPhase, string(corev1.PodSucceeded))
 		pvcPrime.Spec.VolumeName = "test-pv"
 		pvcPrime.UID = pvcPrimeUID
 		err = r.client.Update(context.TODO(), pvcPrime)
@@ -157,14 +156,15 @@ var _ = Describe("Datavolume controller reconcile loop", func() {
 		updatedPVC := &corev1.PersistentVolumeClaim{}
 		err = r.client.Get(context.TODO(), types.NamespacedName{Name: "test-pvc", Namespace: metav1.NamespaceDefault}, updatedPVC)
 		Expect(err).ToNot(HaveOccurred())
-		Expect(updatedPVC.GetAnnotations()).To(BeNil())
+		Expect(updatedPVC.GetAnnotations()[AnnPVCPrimeName]).To(BeEmpty())
 
 		expectEvent(r, uploadSucceeded)
 	})
 
-	It("should clean PVCPrime when targetPVC bound", func() {
+	It("should clean PVCPrime when targetPVC bound and succeeded", func() {
 		pvc := newUploadPopulatorPVC("test-pvc")
 		pvc.Spec.VolumeName = "test-pv"
+		cc.AddAnnotation(pvc, cc.AnnPodPhase, string(corev1.PodSucceeded))
 		pvcPrime := newUploadPopulatorPVC(PVCPrimeName(pvc))
 
 		volumeUploadSourceCR := newUploadPopulatorCR("", false)
@@ -207,8 +207,7 @@ var _ = Describe("Datavolume controller reconcile loop", func() {
 		Expect(err).To(HaveOccurred())
 		Expect(errors.IsNotFound(err)).To(BeTrue())
 
-		pvc.Annotations = make(map[string]string)
-		pvc.Annotations[cc.AnnSelectedNode] = "node01"
+		cc.AddAnnotation(pvc, cc.AnnSelectedNode, "node01")
 		err = r.client.Update(context.TODO(), pvc)
 		Expect(err).ToNot(HaveOccurred())