1 package services
2
3 import (
4 "context"
5 "database/sql"
6 "database/sql/driver"
7 "testing"
8
9 "github.com/DATA-DOG/go-sqlmock"
10 "github.com/golang/mock/gomock"
11 "github.com/stretchr/testify/assert"
12
13 "edge-infra.dev/pkg/edge/api/graph/model"
14 apimock "edge-infra.dev/pkg/edge/api/mocks"
15 "edge-infra.dev/pkg/edge/api/services/artifacts"
16 sqlquery "edge-infra.dev/pkg/edge/api/sql"
17 )
18
19 func TestCreateTerminalDisk(t *testing.T) {
20 db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
21 if err != nil {
22 t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
23 }
24 defer db.Close()
25
26 includeDisk := true
27 expectEmpty := true
28 usePart := false
29 result := sqlmock.NewResult(1, 1)
30
31 mockDBGetTerminalDiskByTerminalID(mock, terminalID, []model.TerminalDisk{})
32
33 mock.ExpectBegin()
34 mockDBTerminalDiskCreate(mock, terminalID, model.TerminalDiskCreateInput{
35 IncludeDisk: includeDisk,
36 ExpectEmpty: expectEmpty,
37 DevicePath: devicePath,
38 UsePart: usePart,
39 }, result)
40 mock.ExpectCommit()
41
42 newDisk := model.TerminalDiskCreateInput{
43 IncludeDisk: includeDisk,
44 ExpectEmpty: expectEmpty,
45 DevicePath: devicePath,
46 UsePart: usePart,
47 }
48 artifactsService := artifacts.NewArtifactsService(db, nil)
49 labelSvc := NewLabelService(artifactsService, db)
50 service := NewTerminalService(db, labelSvc)
51 disk, err := service.CreateTerminalDiskEntry(context.Background(), terminalID, &newDisk)
52 assert.NoError(t, err)
53
54 assert.NotNil(t, disk)
55 assert.Equal(t, disk.IncludeDisk, includeDisk)
56 assert.Equal(t, disk.ExpectEmpty, expectEmpty)
57 assert.Equal(t, devicePath, disk.DevicePath)
58 assert.Equal(t, usePart, disk.UsePart)
59 }
60
61 func TestCreateDuplicateTerminalDisk(t *testing.T) {
62 db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
63 if err != nil {
64 t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
65 }
66 defer db.Close()
67
68 terminalDisk := model.TerminalDisk{
69 TerminalID: terminalID,
70 TerminalDiskID: terminalDiskID,
71 ExpectEmpty: false,
72 IncludeDisk: true,
73 DevicePath: devicePath,
74 UsePart: true,
75 }
76
77 mockDBGetTerminalDiskByTerminalID(mock, terminalID, []model.TerminalDisk{terminalDisk})
78
79 mock.ExpectBegin()
80 mockDBTerminalDiskCreate(mock, terminalID, model.TerminalDiskCreateInput{}, sqlmock.NewResult(1, 0))
81 mock.ExpectRollback()
82
83 mock.ExpectQuery(sqlquery.GetProjectIDByClusterEdgeID).
84 WithArgs(terminalClusterEdgeID).
85 WillReturnRows(mock.NewRows([]string{"project_id"}).
86 AddRow("test-org"))
87
88 getKubeResource := func() GetKubeResourceFunc {
89 return func(_ context.Context, _projectID string, cluster *model.Cluster, _ model.LoqRequest) ([]string, error) {
90 assert.Equal(t, projectID, _projectID)
91 assert.Nil(t, cluster)
92 res := versionResource
93 return []string{res}, nil
94 }
95 }
96 bqClientMock := createMockBQClient(t, getKubeResource())
97 artifactsService := artifacts.NewArtifactsService(db, nil)
98 labelSvc := NewLabelService(artifactsService, db)
99 service := NewTerminalServiceBQ(db, bqClientMock, labelSvc)
100
101 diskInput := model.TerminalDiskCreateInput{
102 IncludeDisk: terminalDisk.IncludeDisk,
103 ExpectEmpty: terminalDisk.ExpectEmpty,
104 DevicePath: terminalDisk.DevicePath,
105 UsePart: terminalDisk.UsePart,
106 }
107
108 _, err = service.CreateTerminalDiskEntry(context.Background(), terminalID, &diskInput)
109 assert.ErrorIs(t, err, ErrDuplicateTerminalDiskDevicePaths)
110 }
111
112 func TestCreateTerminalDiskDefaultNoPart(t *testing.T) {
113 db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
114 if err != nil {
115 t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
116 }
117 defer db.Close()
118
119 includeDisk := true
120 expectEmpty := true
121 result := sqlmock.NewResult(1, 1)
122
123 mockDBGetTerminalDiskByTerminalID(mock, terminalID, []model.TerminalDisk{})
124
125 mock.ExpectBegin()
126 mockDBTerminalDiskCreate(mock, terminalID, model.TerminalDiskCreateInput{
127 IncludeDisk: includeDisk,
128 ExpectEmpty: expectEmpty,
129 DevicePath: devicePath,
130 }, result)
131 mock.ExpectCommit()
132
133 newDisk := model.TerminalDiskCreateInput{
134 IncludeDisk: includeDisk,
135 ExpectEmpty: expectEmpty,
136 DevicePath: devicePath,
137 }
138 artifactsService := artifacts.NewArtifactsService(db, nil)
139 labelSvc := NewLabelService(artifactsService, db)
140 service := NewTerminalService(db, labelSvc)
141 disk, err := service.CreateTerminalDiskEntry(context.Background(), terminalID, &newDisk)
142 assert.NoError(t, err)
143
144 assert.NotNil(t, disk)
145 assert.Equal(t, disk.IncludeDisk, includeDisk)
146 assert.Equal(t, disk.ExpectEmpty, expectEmpty)
147 assert.Equal(t, devicePath, disk.DevicePath)
148 assert.False(t, disk.UsePart)
149 }
150
151 func TestDeleteTerminalDisk(t *testing.T) {
152 db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
153 if err != nil {
154 t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
155 }
156 defer db.Close()
157
158 mockDBGetTerminalIDFromDisk(mock, terminalDiskID, &terminalID)
159
160 mock.ExpectQuery(sqlquery.GetTerminalByIDQuery).
161 WithArgs(terminalID).
162 WillReturnRows(mock.NewRows([]string{"terminal_id", "lane", "role", "cluster_edge_id", "cluster_name", "class", "discover_disks", "boot_disk", "primary_interface", "existing_efi_part", "swap_enabled", "hostname"}).
163 AddRow(terminalID, lane1, terminalRoleWorker, terminalClusterEdgeID, terminalClusterName, terminalClassServer, terminalDiscoverDisksAll, devicePath2, terminalPrimaryInterface, terminalExistingEfiPart, swapEnabled, terminalHostname))
164 mockDBGetTerminalInterfaceByTerminalIDQuery(mock, terminalID, []*model.TerminalInterface{
165 {
166 TerminalInterfaceID: terminalInterfaceID,
167 MacAddress: macAddress,
168 Dhcp4: dhcp4False,
169 Dhcp6: dhcp6False,
170 Gateway4: &gateway4,
171 Gateway6: &gateway6,
172 TerminalID: terminalID,
173 },
174 })
175 mockDBGetTerminalAddressByInterfaceIDQuery(mock, terminalInterfaceID, []*model.TerminalAddress{
176 {
177 TerminalAddressID: terminalAddressID,
178 IP: &ipv4,
179 PrefixLen: prefixLen,
180 Family: familyInet,
181 TerminalInterfaceID: terminalInterfaceID,
182 },
183 })
184 mockDBGetTerminalDiskByTerminalID(mock, terminalID, nil)
185 mock.ExpectQuery(sqlquery.GetTerminalLabels).
186 WithArgs(terminalID, sql.NullString{}).
187 WillReturnRows(mock.NewRows([]string{"terminal_id", "terminal_label_edge_id", "label_edge_id", "labelkey", "color", "visible", "editable", "banner", "unique", "description", "label_type"}).
188 AddRow(terminalID, "388d1144-27c5-44e2-856a-e69a3d4f859f", testLabelEdgeID, label.Key, label.Color, label.Visible, label.Editable, label.BannerEdgeID, label.Unique, label.Description, label.Type))
189 mock.ExpectQuery(sqlquery.GetProjectIDByClusterEdgeID).
190 WithArgs(terminalClusterEdgeID).
191 WillReturnRows(mock.NewRows([]string{"project_id"}).
192 AddRow("test-org"))
193
194 mockDBTerminalDiskDelete(mock, terminalDiskID, sqlmock.NewResult(1, 1))
195
196 ctrl := gomock.NewController(t)
197 defer ctrl.Finish()
198
199 bqmock := apimock.NewMockBQClient(ctrl)
200 bqmock.EXPECT().GetKubeResource(gomock.Any(), "test-org", gomock.Any(), gomock.Any())
201
202 artifactsService := artifacts.NewArtifactsService(db, nil)
203 labelSvc := NewLabelService(artifactsService, db)
204 service := NewTerminalServiceBQ(db, bqmock, labelSvc)
205 newTerminal, err := service.DeleteTerminalDiskEntry(context.Background(), terminalDiskID)
206 assert.NoError(t, err)
207
208 assert.NotNil(t, newTerminal)
209 assert.Equal(t, terminalID, newTerminal.TerminalID)
210 for _, disk := range newTerminal.Disks {
211 assert.NotEqual(t, terminalDiskID, disk.TerminalDiskID)
212 }
213 }
214
215 func TestUpdateTerminalDisk(t *testing.T) {
216 db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
217 if err != nil {
218 t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
219 }
220 defer db.Close()
221
222 disks := []model.TerminalDisk{
223 {
224 TerminalID: terminalID,
225 TerminalDiskID: terminalDiskID,
226 DevicePath: devicePath,
227 ExpectEmpty: true,
228 IncludeDisk: true,
229 UsePart: false,
230 },
231 {
232 TerminalID: terminalID,
233 TerminalDiskID: terminalDiskID2,
234 ExpectEmpty: false,
235 IncludeDisk: true,
236 DevicePath: devicePath2,
237 UsePart: false,
238 },
239 }
240
241 mock.ExpectBegin()
242 mockDBGetTerminalDiskByID(mock, terminalDiskID, &disks[0])
243 mockDBGetTerminalDiskByTerminalID(mock, disks[0].TerminalID, disks)
244
245 newDevicePath := "test-path"
246 mockDBUpdateTerminalDisk(mock, disks[0].TerminalDiskID, model.TerminalDiskUpdateInput{
247 DevicePath: &newDevicePath,
248 IncludeDisk: &disks[0].IncludeDisk,
249 ExpectEmpty: &disks[0].ExpectEmpty,
250 UsePart: &disks[0].UsePart,
251 }, sqlmock.NewResult(1, 1))
252 mock.ExpectCommit()
253
254 ctrl := gomock.NewController(t)
255 defer ctrl.Finish()
256 bqmock := apimock.NewMockBQClient(ctrl)
257
258 artifactsService := artifacts.NewArtifactsService(db, nil)
259 labelSvc := NewLabelService(artifactsService, db)
260 service := NewTerminalServiceBQ(db, bqmock, labelSvc)
261 updatedTerminalDisk, err := service.UpdateTerminalDiskEntry(context.Background(), terminalDiskID, model.TerminalDiskUpdateInput{
262 DevicePath: &newDevicePath,
263 })
264 assert.NoError(t, err)
265 assert.Equal(t, newDevicePath, updatedTerminalDisk.DevicePath)
266 assert.Equal(t, disks[0].ExpectEmpty, updatedTerminalDisk.ExpectEmpty)
267 assert.Equal(t, disks[0].IncludeDisk, updatedTerminalDisk.IncludeDisk)
268 assert.Equal(t, disks[0].UsePart, updatedTerminalDisk.UsePart)
269 }
270
271 func mockDBGetTerminalDiskByID(mock sqlmock.Sqlmock, terminalDiskID string, returnDisk *model.TerminalDisk) {
272 rows := sqlmock.NewRows(terminalDiskColumns)
273 if returnDisk != nil {
274 rows.AddRow(returnDisk.TerminalDiskID, returnDisk.TerminalID, returnDisk.IncludeDisk, returnDisk.ExpectEmpty, returnDisk.DevicePath, returnDisk.UsePart)
275 }
276
277 mock.ExpectQuery(sqlquery.GetTerminalDiskByIDQuery).
278 WithArgs(terminalDiskID).
279 WillReturnRows(rows)
280 }
281
282 func mockDBGetTerminalDiskByTerminalID(mock sqlmock.Sqlmock, terminalID any, returnDisks []model.TerminalDisk) {
283 if terminalID == nil {
284 terminalID = sqlmock.AnyArg()
285 }
286 rows := sqlmock.NewRows(terminalDiskColumns)
287 if len(returnDisks) != 0 {
288 for _, disk := range returnDisks {
289 rows.AddRow(disk.TerminalDiskID, disk.TerminalID, disk.IncludeDisk, disk.ExpectEmpty, disk.DevicePath, disk.UsePart)
290 }
291 }
292
293 mock.ExpectQuery(sqlquery.GetTerminalDiskByTerminalIDQuery).
294 WithArgs(terminalID).
295 WillReturnRows(rows)
296 }
297
298 func mockDBTerminalDiskCreate(mock sqlmock.Sqlmock, terminalID string, diskCreate model.TerminalDiskCreateInput, result driver.Result) {
299 mock.ExpectExec(sqlquery.TerminalDiskCreateQuery).
300 WithArgs(sqlmock.AnyArg(), terminalID, diskCreate.IncludeDisk, diskCreate.ExpectEmpty, diskCreate.DevicePath, diskCreate.UsePart).
301 WillReturnResult(result)
302 }
303
304 func mockDBTerminalDiskDelete(mock sqlmock.Sqlmock, diskID string, result driver.Result) {
305 mock.ExpectExec(sqlquery.TerminalDiskDeleteQuery).
306 WithArgs(diskID).
307 WillReturnResult(result)
308 }
309
310 func mockDBGetTerminalIDFromDisk(mock sqlmock.Sqlmock, diskID string, returnTerminalID *string) {
311 rows := sqlmock.NewRows([]string{"terminal_id"})
312 if returnTerminalID != nil {
313 rows.AddRow(returnTerminalID)
314 }
315
316 mock.ExpectQuery(sqlquery.GetTerminalIDFromDiskQuery).
317 WithArgs(diskID).
318 WillReturnRows(rows)
319 }
320
321 func mockDBUpdateTerminalDisk(mock sqlmock.Sqlmock, terminalDiskID string, updateTerminalDisk model.TerminalDiskUpdateInput, result driver.Result) {
322 mock.ExpectExec(sqlquery.TerminalDiskUpdateQuery).
323 WithArgs(updateTerminalDisk.IncludeDisk, updateTerminalDisk.ExpectEmpty, updateTerminalDisk.DevicePath, updateTerminalDisk.UsePart, terminalDiskID).
324 WillReturnResult(result)
325 }
326
View as plain text