refactor: add builder package to simplify testing code

Signed-off-by: Zespre Schmidt <[email protected]>
starbops · Jan 4, 2025 · 2ddb37a · 2ddb37a
1 parent eea7b76
commit 2ddb37a
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 42 deletions.
diff --git a/pkg/builder/virtualmachine.go b/pkg/builder/virtualmachine.go
@@ -0,0 +1,41 @@
+package builder
+
+import (
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	kubevirtv1 "kubevirt.io/api/core/v1"
+)
+
+// VirtualMachineBuilder builds a VirtualMachine object.
+type VirtualMachineBuilder struct {
+	vm *kubevirtv1.VirtualMachine
+}
+
+func NewVirtualMachineBuilder(namespace, name string) *VirtualMachineBuilder {
+	return &VirtualMachineBuilder{
+		vm: &kubevirtv1.VirtualMachine{
+			ObjectMeta: metav1.ObjectMeta{
+				Namespace: namespace,
+				Name:      name,
+			},
+		},
+	}
+}
+
+func (b *VirtualMachineBuilder) Build() *kubevirtv1.VirtualMachine {
+	return b.vm
+}
+
+func (b *VirtualMachineBuilder) Running(running bool) *VirtualMachineBuilder {
+	b.vm.Spec.Running = &running
+	return b
+}
+
+func (b *VirtualMachineBuilder) RunStrategy(strategy kubevirtv1.VirtualMachineRunStrategy) *VirtualMachineBuilder {
+	b.vm.Spec.RunStrategy = &strategy
+	return b
+}
+
+func (b *VirtualMachineBuilder) Ready(ready bool) *VirtualMachineBuilder {
+	b.vm.Status.Ready = ready
+	return b
+}
diff --git a/pkg/fake/virtualmachine.go b/pkg/fake/virtualmachine.go
@@ -4,7 +4,7 @@ import (
 	"context"
 
 	"github.com/stretchr/testify/mock"
-	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/watch"
 	kubevirtv1 "kubevirt.io/api/core/v1"
@@ -15,37 +15,37 @@ type MockVirtualMachineInterface struct {
 }
 
 func (m *MockVirtualMachineInterface) Get(
-	ctx context.Context, name string, options v1.GetOptions,
+	ctx context.Context, name string, options metav1.GetOptions,
 ) (*kubevirtv1.VirtualMachine, error) {
 	args := m.Called(ctx, name, options)
 	return args.Get(0).(*kubevirtv1.VirtualMachine), args.Error(1)
 }
 
 func (m *MockVirtualMachineInterface) Create(
-	ctx context.Context, vm *kubevirtv1.VirtualMachine, options v1.CreateOptions,
+	ctx context.Context, vm *kubevirtv1.VirtualMachine, options metav1.CreateOptions,
 ) (*kubevirtv1.VirtualMachine, error) {
 	panic("implement me")
 }
 
 func (m *MockVirtualMachineInterface) Update(
-	ctx context.Context, vm *kubevirtv1.VirtualMachine, options v1.UpdateOptions,
+	ctx context.Context, vm *kubevirtv1.VirtualMachine, options metav1.UpdateOptions,
 ) (*kubevirtv1.VirtualMachine, error) {
 	args := m.Called(ctx, vm, options)
 	return args.Get(0).(*kubevirtv1.VirtualMachine), args.Error(1)
 }
 
-func (m *MockVirtualMachineInterface) Delete(ctx context.Context, name string, options v1.DeleteOptions) error {
+func (m *MockVirtualMachineInterface) Delete(ctx context.Context, name string, options metav1.DeleteOptions) error {
 	panic("implement me")
 }
 
 func (m *MockVirtualMachineInterface) DeleteCollection(
-	ctx context.Context, options v1.DeleteOptions, listOptions v1.ListOptions,
+	ctx context.Context, options metav1.DeleteOptions, listOptions metav1.ListOptions,
 ) error {
 	panic("implement me")
 }
 
 func (m *MockVirtualMachineInterface) List(
-	ctx context.Context, options v1.ListOptions,
+	ctx context.Context, options metav1.ListOptions,
 ) (*kubevirtv1.VirtualMachineList, error) {
 	panic("implement me")
 }
@@ -55,18 +55,18 @@ func (m *MockVirtualMachineInterface) Patch(
 	name string,
 	pt types.PatchType,
 	data []byte,
-	options v1.PatchOptions,
+	options metav1.PatchOptions,
 	subresources ...string,
 ) (*kubevirtv1.VirtualMachine, error) {
 	panic("implement me")
 }
 
 func (m *MockVirtualMachineInterface) UpdateStatus(
-	ctx context.Context, vm *kubevirtv1.VirtualMachine, options v1.UpdateOptions,
+	ctx context.Context, vm *kubevirtv1.VirtualMachine, options metav1.UpdateOptions,
 ) (*kubevirtv1.VirtualMachine, error) {
 	panic("implement me")
 }
 
-func (m *MockVirtualMachineInterface) Watch(ctx context.Context, options v1.ListOptions) (watch.Interface, error) {
+func (m *MockVirtualMachineInterface) Watch(ctx context.Context, options metav1.ListOptions) (watch.Interface, error) {
 	panic("implement me")
 }
diff --git a/pkg/fake/virtualmachineinstance.go b/pkg/fake/virtualmachineinstance.go
@@ -4,7 +4,7 @@ import (
 	"context"
 
 	"github.com/stretchr/testify/mock"
-	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/watch"
 	kubevirtv1 "kubevirt.io/api/core/v1"
@@ -15,37 +15,39 @@ type MockVirtualMachineInstanceInterface struct {
 }
 
 func (m *MockVirtualMachineInstanceInterface) Get(
-	ctx context.Context, name string, options v1.GetOptions,
+	ctx context.Context, name string, options metav1.GetOptions,
 ) (*kubevirtv1.VirtualMachineInstance, error) {
 	panic("implement me")
 }
 
 func (m *MockVirtualMachineInstanceInterface) Create(
-	ctx context.Context, vmi *kubevirtv1.VirtualMachineInstance, options v1.CreateOptions,
+	ctx context.Context, vmi *kubevirtv1.VirtualMachineInstance, options metav1.CreateOptions,
 ) (*kubevirtv1.VirtualMachineInstance, error) {
 	panic("implement me")
 }
 
 func (m *MockVirtualMachineInstanceInterface) Update(
-	ctx context.Context, vmi *kubevirtv1.VirtualMachineInstance, options v1.UpdateOptions,
+	ctx context.Context, vmi *kubevirtv1.VirtualMachineInstance, options metav1.UpdateOptions,
 ) (*kubevirtv1.VirtualMachineInstance, error) {
 	args := m.Called(ctx, vmi, options)
 	return args.Get(0).(*kubevirtv1.VirtualMachineInstance), args.Error(1)
 }
 
-func (m *MockVirtualMachineInstanceInterface) Delete(ctx context.Context, name string, options v1.DeleteOptions) error {
+func (m *MockVirtualMachineInstanceInterface) Delete(
+	ctx context.Context, name string, options metav1.DeleteOptions,
+) error {
 	args := m.Called(ctx, name, options)
 	return args.Error(0)
 }
 
 func (m *MockVirtualMachineInstanceInterface) DeleteCollection(
-	ctx context.Context, options v1.DeleteOptions, listOptions v1.ListOptions,
+	ctx context.Context, options metav1.DeleteOptions, listOptions metav1.ListOptions,
 ) error {
 	panic("implement me")
 }
 
 func (m *MockVirtualMachineInstanceInterface) List(
-	ctx context.Context, options v1.ListOptions,
+	ctx context.Context, options metav1.ListOptions,
 ) (*kubevirtv1.VirtualMachineInstanceList, error) {
 	panic("implement me")
 }
@@ -55,20 +57,20 @@ func (m *MockVirtualMachineInstanceInterface) Patch(
 	name string,
 	pt types.PatchType,
 	data []byte,
-	options v1.PatchOptions,
+	options metav1.PatchOptions,
 	subresources ...string,
 ) (*kubevirtv1.VirtualMachineInstance, error) {
 	panic("implement me")
 }
 
 func (m *MockVirtualMachineInstanceInterface) UpdateStatus(
-	ctx context.Context, vm *kubevirtv1.VirtualMachineInstance, options v1.UpdateOptions,
+	ctx context.Context, vm *kubevirtv1.VirtualMachineInstance, options metav1.UpdateOptions,
 ) (*kubevirtv1.VirtualMachineInstance, error) {
 	panic("implement me")
 }
 
 func (m *MockVirtualMachineInstanceInterface) Watch(
-	ctx context.Context, options v1.ListOptions,
+	ctx context.Context, options metav1.ListOptions,
 ) (watch.Interface, error) {
 	panic("implement me")
 }
diff --git a/pkg/resourcemanager/virtualmachine_test.go b/pkg/resourcemanager/virtualmachine_test.go
@@ -8,6 +8,7 @@ import (
 	"github.com/stretchr/testify/require"
 	kubevirtv1 "kubevirt.io/api/core/v1"
 
+	"kubevirt.io/kubevirtbmc/pkg/builder"
 	"kubevirt.io/kubevirtbmc/pkg/fake"
 )
 
@@ -16,11 +17,8 @@ func TestGetPowerStatus(t *testing.T) {
 	mockVMInterface := new(fake.MockVirtualMachineInterface)
 	mockClient.On("VirtualMachines", "default").Return(mockVMInterface)
 
-	mockVM := &kubevirtv1.VirtualMachine{
-		Status: kubevirtv1.VirtualMachineStatus{
-			Ready: true,
-		},
-	}
+	mockVM := builder.NewVirtualMachineBuilder("default", "test-vm").Ready(true).Build()
+
 	mockVMInterface.On("Get", mock.Anything, "test-vm", mock.Anything).Return(mockVM, nil)
 
 	vmrm := &VirtualMachineResourceManager{
@@ -37,7 +35,6 @@ func TestGetPowerStatus(t *testing.T) {
 
 	// Add another test case where the VM is not ready
 	mockVM.Status.Ready = false
-	mockVMInterface.On("Get", mock.Anything, "test-vm", mock.Anything).Return(mockVM, nil)
 
 	status, err = vmrm.GetPowerStatus()
 	require.NoError(t, err)
@@ -49,13 +46,11 @@ func TestPowerOn(t *testing.T) {
 	mockVMInterface := new(fake.MockVirtualMachineInterface)
 	mockClient.On("VirtualMachines", "default").Return(mockVMInterface)
 
-	mockVM := &kubevirtv1.VirtualMachine{
-		Spec: kubevirtv1.VirtualMachineSpec{
-			Running: func(b bool) *bool { return &b }(false),
-		},
-	}
-	mockVMInterface.On("Get", mock.Anything, "test-vm", mock.Anything).Return(mockVM, nil)
-	mockVMInterface.On("Update", mock.Anything, mockVM, mock.Anything).Return(mockVM, nil)
+	mockVM := builder.NewVirtualMachineBuilder("default", "test-vm").Running(false).Build()
+
+	mockVMInterface.
+		On("Get", mock.Anything, "test-vm", mock.Anything).Return(mockVM, nil).
+		On("Update", mock.Anything, mockVM, mock.Anything).Return(mockVM, nil)
 
 	vmrm := &VirtualMachineResourceManager{
 		ctx:       context.TODO(),
@@ -81,13 +76,11 @@ func TestPowerOff(t *testing.T) {
 	mockVMInterface := new(fake.MockVirtualMachineInterface)
 	mockClient.On("VirtualMachines", "default").Return(mockVMInterface)
 
-	mockVM := &kubevirtv1.VirtualMachine{
-		Spec: kubevirtv1.VirtualMachineSpec{
-			Running: func(b bool) *bool { return &b }(false),
-		},
-	}
-	mockVMInterface.On("Get", mock.Anything, "test-vm", mock.Anything).Return(mockVM, nil)
-	mockVMInterface.On("Update", mock.Anything, mockVM, mock.Anything).Return(mockVM, nil)
+	mockVM := builder.NewVirtualMachineBuilder("default", "test-vm").Running(true).Build()
+
+	mockVMInterface.
+		On("Get", mock.Anything, "test-vm", mock.Anything).Return(mockVM, nil).
+		On("Update", mock.Anything, mockVM, mock.Anything).Return(mockVM, nil)
 
 	vmrm := &VirtualMachineResourceManager{
 		ctx:       context.TODO(),
@@ -159,8 +152,10 @@ func TestSetBootDevice(t *testing.T) {
 			},
 		},
 	}
-	mockVMInterface.On("Get", mock.Anything, "test-vm", mock.Anything).Return(mockVM, nil)
-	mockVMInterface.On("Update", mock.Anything, mockVM, mock.Anything).Return(mockVM, nil)
+
+	mockVMInterface.
+		On("Get", mock.Anything, "test-vm", mock.Anything).Return(mockVM, nil).
+		On("Update", mock.Anything, mockVM, mock.Anything).Return(mockVM, nil)
 
 	vmrm := &VirtualMachineResourceManager{
 		ctx:       context.TODO(),