Use DV options instead of params

Signed-off-by: Arnon Gilboa <[email protected]>
kiagnose · Aug 12, 2024 · ead00c1 · ead00c1
1 parent 9701e5e
commit ead00c1
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 24 deletions.
diff --git a/pkg/internal/checkup/vmi.go b/pkg/internal/checkup/vmi.go
@@ -40,16 +40,34 @@ const (
 func newVMUnderTest(name string, pvc *corev1.PersistentVolumeClaim, snap *snapshotv1.VolumeSnapshot,
 	checkupConfig config.Config, addBlankDataVolume bool) *kvcorev1.VirtualMachine {
 	dvName := getVMDvName(name)
+	dvOpts := []vmi.DataVolumeOption{}
+
+	if pvc != nil {
+		dvOpts = append(dvOpts, vmi.WithDataVolumePvcSource(pvc))
+	} else if snap != nil {
+		dvOpts = append(dvOpts, vmi.WithDataVolumeSnapshotSource(snap))
+	}
+
+	if checkupConfig.StorageClass != "" {
+		dvOpts = append(dvOpts, vmi.WithDataVolumeStorageClass(checkupConfig.StorageClass))
+	}
+
 	optionsToApply := []vmi.Option{
-		vmi.WithDataVolume(dvName, pvc, snap, checkupConfig.StorageClass),
+		vmi.WithDataVolume(dvName, dvOpts...),
 		vmi.WithMemory(guestMemory),
 		vmi.WithTerminationGracePeriodSeconds(terminationGracePeriodSeconds),
 		vmi.WithOwnerReference(checkupConfig.PodName, checkupConfig.PodUID),
 	}
+
 	if addBlankDataVolume {
 		blankDvName := fmt.Sprintf("%s-blank", dvName)
-		optionsToApply = append(optionsToApply, vmi.WithDataVolume(blankDvName, nil, nil, checkupConfig.StorageClass))
+		dvOpts := []vmi.DataVolumeOption{vmi.WithDataVolumeBlankSource()}
+		if checkupConfig.StorageClass != "" {
+			dvOpts = append(dvOpts, vmi.WithDataVolumeStorageClass(checkupConfig.StorageClass))
+		}
+		optionsToApply = append(optionsToApply, vmi.WithDataVolume(blankDvName, dvOpts...))
 	}
+
 	return vmi.NewVM(name, optionsToApply...)
 }
 

diff --git a/pkg/internal/checkup/vmi/spec.go b/pkg/internal/checkup/vmi/spec.go
@@ -54,7 +54,9 @@ func NewVM(name string, options ...Option) *kvcorev1.VirtualMachine {
 	return newVM
 }
 
-func WithDataVolume(dvName string, pvc *corev1.PersistentVolumeClaim, snap *snapshotv1.VolumeSnapshot, storageClass string) Option {
+type DataVolumeOption func(*cdiv1.DataVolumeSpec)
+
+func WithDataVolume(dvName string, opts ...DataVolumeOption) Option {
 	return func(vm *kvcorev1.VirtualMachine) {
 		dvt := kvcorev1.DataVolumeTemplateSpec{
 			ObjectMeta: metav1.ObjectMeta{
@@ -66,27 +68,8 @@ func WithDataVolume(dvName string, pvc *corev1.PersistentVolumeClaim, snap *snap
 			},
 		}
 
-		if pvc != nil {
-			dvt.Spec.Source.PVC = &cdiv1.DataVolumeSourcePVC{
-				Namespace: pvc.Namespace,
-				Name:      pvc.Name,
-			}
-		} else if snap != nil {
-			dvt.Spec.Source.Snapshot = &cdiv1.DataVolumeSourceSnapshot{
-				Namespace: snap.Namespace,
-				Name:      snap.Name,
-			}
-		} else {
-			dvt.Spec.Source.Blank = &cdiv1.DataVolumeBlankImage{}
-			dvt.Spec.Storage.Resources.Requests = corev1.ResourceList{
-				corev1.ResourceStorage: resource.MustParse("1G"),
-			}
-		}
-
-		if storageClass != "" {
-			dvt.Spec.Storage.StorageClassName = &storageClass
-		} else if pvc != nil {
-			dvt.Spec.Storage.StorageClassName = pvc.Spec.StorageClassName
+		for _, f := range opts {
+			f(&dvt.Spec)
 		}
 
 		vm.Spec.DataVolumeTemplates = append(vm.Spec.DataVolumeTemplates, dvt)
@@ -103,6 +86,43 @@ func WithDataVolume(dvName string, pvc *corev1.PersistentVolumeClaim, snap *snap
 	}
 }
 
+func WithDataVolumePvcSource(pvc *corev1.PersistentVolumeClaim) DataVolumeOption {
+	return func(dvSpec *cdiv1.DataVolumeSpec) {
+		dvSpec.Source.PVC = &cdiv1.DataVolumeSourcePVC{
+			Namespace: pvc.Namespace,
+			Name:      pvc.Name,
+		}
+
+		if dvSpec.Storage.StorageClassName == nil {
+			dvSpec.Storage.StorageClassName = pvc.Spec.StorageClassName
+		}
+	}
+}
+
+func WithDataVolumeSnapshotSource(snap *snapshotv1.VolumeSnapshot) DataVolumeOption {
+	return func(dvSpec *cdiv1.DataVolumeSpec) {
+		dvSpec.Source.Snapshot = &cdiv1.DataVolumeSourceSnapshot{
+			Namespace: snap.Namespace,
+			Name:      snap.Name,
+		}
+	}
+}
+
+func WithDataVolumeBlankSource() DataVolumeOption {
+	return func(dvSpec *cdiv1.DataVolumeSpec) {
+		dvSpec.Source.Blank = &cdiv1.DataVolumeBlankImage{}
+		dvSpec.Storage.Resources.Requests = corev1.ResourceList{
+			corev1.ResourceStorage: resource.MustParse("1Gi"),
+		}
+	}
+}
+
+func WithDataVolumeStorageClass(storageClass string) DataVolumeOption {
+	return func(dvSpec *cdiv1.DataVolumeSpec) {
+		dvSpec.Storage.StorageClassName = &storageClass
+	}
+}
+
 func WithMemory(guestMemory string) Option {
 	return func(vm *kvcorev1.VirtualMachine) {
 		guestMemoryQuantity := resource.MustParse(guestMemory)