1 package services
2
3 import (
4 "context"
5 "database/sql/driver"
6 "testing"
7
8 "github.com/DATA-DOG/go-sqlmock"
9 "github.com/stretchr/testify/assert"
10
11 "edge-infra.dev/pkg/edge/api/graph/model"
12 sqlquery "edge-infra.dev/pkg/edge/api/sql"
13 )
14
15 func TestCreateVirtualMachineDisk(t *testing.T) {
16 mockDisk := virtualMachines[0].Disks[0]
17
18 createErr, mockErr := genericTestCreateVirtualMachineDisk(t, *mockDisk)
19
20 assert.NoError(t, createErr)
21 assert.NoError(t, mockErr)
22 }
23
24 func TestCreateVirtualMachineDiskNilContainerImageURL(t *testing.T) {
25 mockDisk := virtualMachines[1].Disks[0]
26
27 createErr, mockErr := genericTestCreateVirtualMachineDisk(t, *mockDisk)
28
29 assert.NoError(t, createErr)
30 assert.NoError(t, mockErr)
31 }
32
33 func TestDeleteVirtualMachineDisk(t *testing.T) {
34 db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
35 if err != nil {
36 t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
37 }
38 defer db.Close()
39
40 mockVM := virtualMachines[0]
41 virtualMachineID := mockVM.VirtualMachineID
42 diskID := mockVM.Disks[0].DiskID
43
44
45 mockDBGetVirtualMachineIDFromDisk(mock, diskID, &virtualMachineID)
46 mockDBGetVirtualMachineByID(mock, virtualMachineID, &mockVM)
47
48 mockDBGetVirtualMachineByID(mock, virtualMachineID, &mockVM)
49 mockDBGetVirtualMachineDisks(mock, []*model.VirtualMachineDisk{
50 mockVM.Disks[0],
51 mockVM.Disks[1],
52 })
53 mockDBVirtualMachineDiskDelete(mock, diskID, sqlmock.NewResult(1, 1))
54
55
56 mockDBGetVirtualMachineByID(mock, virtualMachineID, &mockVM)
57 mockDBGetVirtualMachineDisks(mock, []*model.VirtualMachineDisk{
58 mockVM.Disks[1],
59 })
60
61 service := NewVirtualMachineService(db)
62 ctx := context.Background()
63 _, err = service.DeleteVirtualMachineDiskEntry(ctx, diskID)
64 assert.NoError(t, err)
65
66 updatedDisks, err := service.GetVirtualMachineDisks(ctx, virtualMachineID)
67 assert.NoError(t, err)
68
69 for _, disk := range updatedDisks {
70 assert.NotEqual(t, diskID, disk.DiskID)
71 }
72 assert.NoError(t, mock.ExpectationsWereMet())
73 }
74
75 func TestUpdateVirtualMachineDisk(t *testing.T) {
76 db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
77 if err != nil {
78 t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
79 }
80 defer db.Close()
81
82 disks := []model.VirtualMachineDisk{*virtualMachines[0].Disks[0], *virtualMachines[0].Disks[1]}
83
84 mock.ExpectBegin()
85 mockDBGetVirtualMachineDiskByID(mock, disks[0].DiskID, &disks[0])
86
87 newSize := "128Gi"
88 newContainerImageURL := "docker://us-east1-docker.pkg.dev/ret-edge-pltf-infra/workloads/ien-image:v1.7.0"
89 mockDBUpdateVirtualMachineDisk(mock, disks[0].DiskID, model.VirtualMachineDiskUpdateInput{
90 Type: &disks[0].Type,
91 Bus: &disks[0].Bus,
92 BootOrder: &disks[0].BootOrder,
93 Size: &newSize,
94 ContainerImageURL: &newContainerImageURL,
95 }, sqlmock.NewResult(1, 1))
96 mockDBGetVirtualMachineDisks(mock, []*model.VirtualMachineDisk{
97 &disks[0],
98 &disks[1],
99 })
100 mock.ExpectCommit()
101
102 service := NewVirtualMachineService(db)
103 updateDisk := model.VirtualMachineDiskIDInput{
104 DiskID: disks[0].DiskID,
105 VirtualMachineDiskValues: &model.VirtualMachineDiskUpdateInput{
106 Size: &newSize,
107 ContainerImageURL: &newContainerImageURL,
108 },
109 }
110 updateDisks := []*model.VirtualMachineDiskIDInput{&updateDisk}
111 updatedDisks, err := service.UpdateVirtualMachineDiskEntries(context.Background(), updateDisks)
112 assert.NoError(t, err)
113 assert.NoError(t, mock.ExpectationsWereMet())
114 updatedDisk := updatedDisks[0]
115 assert.Equal(t, disks[0].Type, updatedDisk.Type)
116 assert.Equal(t, disks[0].Bus, updatedDisk.Bus)
117 assert.Equal(t, disks[0].BootOrder, updatedDisk.BootOrder)
118 assert.Equal(t, newSize, updatedDisk.Size)
119 assert.Equal(t, newContainerImageURL, updatedDisk.ContainerImageURL)
120 }
121
122 func mockDBVirtualMachineDiskCreate(mock sqlmock.Sqlmock, virtualMachineID string, createDisk model.VirtualMachineDiskCreateInput, result driver.Result) {
123 mock.ExpectExec(sqlquery.VirtualMachineDiskCreateQuery).
124 WithArgs(
125 sqlmock.AnyArg(),
126 virtualMachineID,
127 createDisk.Type,
128 createDisk.Bus,
129 createDisk.BootOrder,
130 createDisk.Size,
131 createDisk.ContainerImageURL).
132 WillReturnResult(result)
133 }
134
135 func mockDBVirtualMachineDisksCreate(mock sqlmock.Sqlmock, createDisks []*model.VirtualMachineDiskCreateInput, result driver.Result) {
136 for _, createDisk := range createDisks {
137 mock.ExpectExec(sqlquery.VirtualMachineDiskCreateQuery).
138 WithArgs(
139 sqlmock.AnyArg(),
140 sqlmock.AnyArg(),
141 createDisk.Type,
142 createDisk.Bus,
143 createDisk.BootOrder,
144 createDisk.Size,
145 createDisk.ContainerImageURL).
146 WillReturnResult(result)
147 }
148 }
149
150 func mockDBVirtualMachineDiskDelete(mock sqlmock.Sqlmock, diskID string, result driver.Result) {
151 mock.ExpectExec(sqlquery.VirtualMachineDiskDeleteQuery).
152 WithArgs(diskID).
153 WillReturnResult(result)
154 }
155
156 func mockDBGetVirtualMachineIDFromDisk(mock sqlmock.Sqlmock, diskID string, returnVirtualMachineID *string) {
157 rows := sqlmock.NewRows([]string{"vm_id"})
158 if returnVirtualMachineID != nil {
159 rows.AddRow(returnVirtualMachineID)
160 }
161
162 mock.ExpectQuery(sqlquery.GetVirtualMachineIDFromDiskQuery).
163 WithArgs(diskID).
164 WillReturnRows(rows)
165 }
166
167 func mockDBUpdateVirtualMachineDisk(mock sqlmock.Sqlmock, diskID string, updateDisk model.VirtualMachineDiskUpdateInput, result driver.Result) {
168 mock.ExpectExec(sqlquery.VirtualMachineDiskUpdateQuery).
169 WithArgs(updateDisk.Type, updateDisk.Bus, updateDisk.BootOrder, updateDisk.Size, updateDisk.ContainerImageURL, diskID).
170 WillReturnResult(result)
171 }
172
173 func mockDBGetVirtualMachineDiskByID(mock sqlmock.Sqlmock, diskID string, returnDisk *model.VirtualMachineDisk) {
174 rows := sqlmock.NewRows(virtualMachineDiskColumns)
175 if returnDisk != nil {
176 rows.AddRow(returnDisk.DiskID, returnDisk.VirtualMachineID, returnDisk.Type, returnDisk.Bus, returnDisk.BootOrder, returnDisk.Size, returnDisk.ContainerImageURL)
177 }
178
179 mock.ExpectQuery(sqlquery.GetVirtualMachineDiskByIDQuery).
180 WithArgs(diskID).
181 WillReturnRows(rows)
182 }
183
184 func mockDBGetVirtualMachineDisks(mock sqlmock.Sqlmock, returnDisks []*model.VirtualMachineDisk) {
185 rows := mock.NewRows(virtualMachineDiskColumns)
186
187 if len(returnDisks) != 0 {
188 for _, returnDisk := range returnDisks {
189 rows.AddRow(returnDisk.DiskID, returnDisk.VirtualMachineID, returnDisk.Type, returnDisk.Bus, returnDisk.BootOrder, returnDisk.Size, returnDisk.ContainerImageURL)
190 }
191 }
192
193 mock.ExpectQuery(sqlquery.GetVirtualMachineDiskByVirtualMachineIDQuery).
194 WillReturnRows(rows)
195 }
196
197 func getVirtualMachineDiskCreateInput(disk model.VirtualMachineDisk) *model.VirtualMachineDiskCreateInput {
198 return &model.VirtualMachineDiskCreateInput{
199 Type: &disk.Type,
200 Bus: &disk.Bus,
201 BootOrder: disk.BootOrder,
202 Size: &disk.Size,
203 ContainerImageURL: &disk.ContainerImageURL,
204 }
205 }
206
207 func genericTestCreateVirtualMachineDisk(t *testing.T, mockDisk model.VirtualMachineDisk) (error, error) {
208 db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
209 if err != nil {
210 t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
211 }
212 defer db.Close()
213 createDisk := getVirtualMachineDiskCreateInput(mockDisk)
214 mock.ExpectBegin()
215 mockDBVirtualMachineDiskCreate(mock, mockDisk.VirtualMachineID, *createDisk, sqlmock.NewResult(1, 1))
216 mockDBGetVirtualMachineDisks(mock, []*model.VirtualMachineDisk{
217 &mockDisk,
218 })
219 mock.ExpectCommit()
220
221 service := NewVirtualMachineService(db)
222
223 createDisks := []*model.VirtualMachineDiskCreateInput{createDisk}
224 _, createErr := service.CreateVirtualMachineDiskEntries(context.Background(), mockDisk.VirtualMachineID, createDisks)
225
226 mockErr := mock.ExpectationsWereMet()
227
228 return createErr, mockErr
229 }
230
View as plain text