Dragonfly2

1
/*
2
 *     Copyright 2023 The Dragonfly Authors
3
 *
4
 * Licensed under the Apache License, Version 2.0 (the "License");
5
 * you may not use this file except in compliance with the License.
6
 * You may obtain a copy of the License at
7
 *
8
 *      http://www.apache.org/licenses/LICENSE-2.0
9
 *
10
 * Unless required by applicable law or agreed to in writing, software
11
 * distributed under the License is distributed on an "AS IS" BASIS,
12
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
 * See the License for the specific language governing permissions and
14
 * limitations under the License.
15
 */
16

17
package storage
18

19
import (
20
	"fmt"
21
	"io/fs"
22
	"os"
23
	"path/filepath"
24
	"reflect"
25
	"regexp"
26
	"testing"
27

28
	"github.com/gocarina/gocsv"
29
	"github.com/stretchr/testify/assert"
30
	"github.com/stretchr/testify/require"
31

32
	schedulerstorage "d7y.io/dragonfly/v2/scheduler/storage"
33
)
34

35
var mockModelKey = "bar"
36

37
func TestStorage_New(t *testing.T) {
38
	tests := []struct {
39
		name    string
40
		baseDir string
41
		expect  func(t *testing.T, s Storage)
42
	}{
43
		{
44
			name:    "new storage",
45
			baseDir: os.TempDir(),
46
			expect: func(t *testing.T, s Storage) {
47
				assert := assert.New(t)
48
				assert.Equal(reflect.TypeOf(s).Elem().Name(), "storage")
49
			},
50
		},
51
	}
52

53
	for _, tc := range tests {
54
		t.Run(tc.name, func(t *testing.T) {
55
			tc.expect(t, New(tc.baseDir))
56
		})
57
	}
58
}
59

60
func TestStorage_ListDownload(t *testing.T) {
61
	require := require.New(t)
62
	testData, err := os.ReadFile("./testdata/download.csv")
63
	require.Nil(err, "load test file")
64

65
	tests := []struct {
66
		name    string
67
		baseDir string
68
		mock    func(t *testing.T, s Storage, baseDir, modelKey string, download []byte)
69
		expect  func(t *testing.T, s Storage, baseDir, modelKey string, download []byte)
70
	}{
71
		{
72
			name:    "empty csv file given",
73
			baseDir: os.TempDir(),
74
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
75
				file, err := os.OpenFile(filepath.Join(baseDir, "download_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
76
				if err != nil {
77
					t.Fatal(err)
78
				}
79
				defer file.Close()
80
			},
81
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
82
				assert := assert.New(t)
83
				_, err := s.ListDownload(modelKey)
84
				assert.EqualError(err, "empty csv file given")
85
			},
86
		},
87
		{
88
			name:    "get file failed",
89
			baseDir: os.TempDir(),
90
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
91
				file, err := os.OpenFile(filepath.Join(baseDir, "download_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
92
				if err != nil {
93
					t.Fatal(err)
94
				}
95
				defer file.Close()
96

97
				if _, err = file.Write(download); err != nil {
98
					t.Fatal(err)
99
				}
100
				s.(*storage).baseDir = "bas"
101
			},
102
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
103
				assert := assert.New(t)
104
				_, err := s.ListDownload(modelKey)
105
				assert.EqualError(err, "open bas/download_bar.csv: no such file or directory")
106
				s.(*storage).baseDir = baseDir
107
			},
108
		},
109
		{
110
			name:    "list downloads of a file",
111
			baseDir: os.TempDir(),
112
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
113
				file, err := os.OpenFile(filepath.Join(baseDir, "download_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
114
				if err != nil {
115
					t.Fatal(err)
116
				}
117
				defer file.Close()
118

119
				if _, err = file.Write(download); err != nil {
120
					t.Fatal(err)
121
				}
122
			},
123
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
124
				assert := assert.New(t)
125
				list, err := s.ListDownload(modelKey)
126
				assert.NoError(err)
127
				assert.Equal(len(list), 1)
128
			},
129
		},
130
	}
131

132
	for _, tc := range tests {
133
		t.Run(tc.name, func(t *testing.T) {
134
			s := New(tc.baseDir)
135
			tc.mock(t, s, tc.baseDir, mockModelKey, testData)
136
			tc.expect(t, s, tc.baseDir, mockModelKey, testData)
137
			if err := s.ClearDownload(mockModelKey); err != nil {
138
				t.Fatal(err)
139
			}
140
		})
141
	}
142
}
143

144
func TestStorage_ListNetworkTopology(t *testing.T) {
145
	require := require.New(t)
146
	testData, err := os.ReadFile("./testdata/networktopology.csv")
147
	require.Nil(err, "load test file")
148

149
	tests := []struct {
150
		name    string
151
		baseDir string
152
		mock    func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte)
153
		expect  func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte)
154
	}{
155
		{
156
			name:    "empty csv file given",
157
			baseDir: os.TempDir(),
158
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
159
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
160
				if err != nil {
161
					t.Fatal(err)
162
				}
163
				defer file.Close()
164
			},
165
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
166
				assert := assert.New(t)
167
				_, err := s.ListNetworkTopology(modelKey)
168
				assert.EqualError(err, "empty csv file given")
169
			},
170
		},
171
		{
172
			name:    "get file failed",
173
			baseDir: os.TempDir(),
174
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
175
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
176
				if err != nil {
177
					t.Fatal(err)
178
				}
179
				defer file.Close()
180

181
				if _, err = file.Write(networkTopology); err != nil {
182
					t.Fatal(err)
183
				}
184
				s.(*storage).baseDir = "foo"
185
			},
186
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
187
				assert := assert.New(t)
188
				_, err := s.ListNetworkTopology(modelKey)
189
				assert.EqualError(err, "open foo/networktopology_bar.csv: no such file or directory")
190
				s.(*storage).baseDir = baseDir
191
			},
192
		},
193
		{
194
			name:    "list network topologies of a file",
195
			baseDir: os.TempDir(),
196
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
197
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
198
				if err != nil {
199
					t.Fatal(err)
200
				}
201
				defer file.Close()
202

203
				if _, err = file.Write(networkTopology); err != nil {
204
					t.Fatal(err)
205
				}
206
			},
207
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
208
				assert := assert.New(t)
209
				list, err := s.ListNetworkTopology(modelKey)
210
				assert.NoError(err)
211
				assert.Equal(len(list), 1)
212
			},
213
		},
214
	}
215

216
	for _, tc := range tests {
217
		t.Run(tc.name, func(t *testing.T) {
218
			s := New(tc.baseDir)
219
			tc.mock(t, s, tc.baseDir, mockModelKey, testData)
220
			tc.expect(t, s, tc.baseDir, mockModelKey, testData)
221
			if err := s.ClearNetworkTopology(mockModelKey); err != nil {
222
				t.Fatal(err)
223
			}
224
		})
225
	}
226
}
227

228
func TestStorage_OpenDownload(t *testing.T) {
229
	require := require.New(t)
230
	testData, err := os.ReadFile("./testdata/download.csv")
231
	require.Nil(err, "load test file")
232

233
	tests := []struct {
234
		name    string
235
		baseDir string
236
		mock    func(t *testing.T, s Storage, baseDir, modelKey string, download []byte)
237
		expect  func(t *testing.T, s Storage, baseDir, modelKey string, download []byte)
238
	}{
239
		{
240
			name:    "open file failed",
241
			baseDir: os.TempDir(),
242
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
243
				file, err := os.OpenFile(filepath.Join(baseDir, "download_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
244
				if err != nil {
245
					t.Fatal(err)
246
				}
247
				defer file.Close()
248
				s.(*storage).baseDir = "baw"
249
			},
250
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
251
				assert := assert.New(t)
252
				_, err := s.OpenDownload(modelKey)
253
				assert.EqualError(err, "open baw/download_bar.csv: no such file or directory")
254
				s.(*storage).baseDir = baseDir
255
			},
256
		},
257
		{
258
			name:    "open storage with downloads of a file",
259
			baseDir: os.TempDir(),
260
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
261
				file, err := os.OpenFile(filepath.Join(baseDir, "download_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
262
				if err != nil {
263
					t.Fatal(err)
264
				}
265
				defer file.Close()
266

267
				if _, err = file.Write(download); err != nil {
268
					t.Fatal(err)
269
				}
270
			},
271
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
272
				assert := assert.New(t)
273
				readCloser, err := s.OpenDownload(modelKey)
274
				assert.NoError(err)
275

276
				var downloads []schedulerstorage.Download
277
				assert.NoError(gocsv.UnmarshalWithoutHeaders(readCloser, &downloads))
278
				assert.Equal(len(downloads), 1)
279
			},
280
		},
281
	}
282

283
	for _, tc := range tests {
284
		t.Run(tc.name, func(t *testing.T) {
285
			s := New(tc.baseDir)
286
			tc.mock(t, s, tc.baseDir, mockModelKey, testData)
287
			tc.expect(t, s, tc.baseDir, mockModelKey, testData)
288
			if err := s.ClearDownload(mockModelKey); err != nil {
289
				t.Fatal(err)
290
			}
291
		})
292
	}
293
}
294

295
func TestStorage_OpenNetworkTopology(t *testing.T) {
296
	require := require.New(t)
297
	testData, err := os.ReadFile("./testdata/networktopology.csv")
298
	require.Nil(err, "load test file")
299

300
	tests := []struct {
301
		name    string
302
		baseDir string
303
		mock    func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte)
304
		expect  func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte)
305
	}{
306
		{
307
			name:    "open file failed",
308
			baseDir: os.TempDir(),
309
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
310
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
311
				if err != nil {
312
					t.Fatal(err)
313
				}
314
				defer file.Close()
315
				s.(*storage).baseDir = "bas"
316
			},
317
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
318
				assert := assert.New(t)
319
				_, err := s.OpenNetworkTopology(modelKey)
320
				assert.EqualError(err, "open bas/networktopology_bar.csv: no such file or directory")
321
				s.(*storage).baseDir = baseDir
322
			},
323
		},
324
		{
325
			name:    "open storage with network topologies of a file",
326
			baseDir: os.TempDir(),
327
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
328
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
329
				if err != nil {
330
					t.Fatal(err)
331
				}
332
				defer file.Close()
333

334
				if _, err = file.Write(networkTopology); err != nil {
335
					t.Fatal(err)
336
				}
337
			},
338
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
339
				assert := assert.New(t)
340
				readCloser, err := s.OpenNetworkTopology(modelKey)
341
				assert.NoError(err)
342

343
				var networkTopologies []schedulerstorage.NetworkTopology
344
				assert.NoError(gocsv.UnmarshalWithoutHeaders(readCloser, &networkTopologies))
345
				assert.Equal(len(networkTopologies), 1)
346
			},
347
		},
348
	}
349

350
	for _, tc := range tests {
351
		t.Run(tc.name, func(t *testing.T) {
352
			s := New(tc.baseDir)
353
			tc.mock(t, s, tc.baseDir, mockModelKey, testData)
354
			tc.expect(t, s, tc.baseDir, mockModelKey, testData)
355
			if err := s.ClearNetworkTopology(mockModelKey); err != nil {
356
				t.Fatal(err)
357
			}
358
		})
359
	}
360
}
361

362
func TestStorage_ClearDownload(t *testing.T) {
363
	tests := []struct {
364
		name    string
365
		baseDir string
366
		mock    func(t *testing.T, s Storage, baseDir, modelKey string)
367
		expect  func(t *testing.T, s Storage, baseDir, modelKey string)
368
	}{
369
		{
370
			name:    "clear file",
371
			baseDir: os.TempDir(),
372
			mock: func(t *testing.T, s Storage, baseDir, modelKey string) {
373
				file, err := os.OpenFile(filepath.Join(baseDir, "download_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
374
				if err != nil {
375
					t.Fatal(err)
376
				}
377
				defer file.Close()
378
			},
379
			expect: func(t *testing.T, s Storage, baseDir, modelKey string) {
380
				assert := assert.New(t)
381
				assert.NoError(s.ClearDownload(modelKey))
382
				fileInfos, err := os.ReadDir(filepath.Join(baseDir))
383
				assert.NoError(err)
384

385
				var backups []fs.FileInfo
386
				re := regexp.MustCompile(fmt.Sprintf("%s_%s", DownloadFilePrefix, modelKey))
387

388
				for _, fileInfo := range fileInfos {
389
					if !fileInfo.IsDir() && re.MatchString(fileInfo.Name()) {
390
						info, _ := fileInfo.Info()
391
						backups = append(backups, info)
392
					}
393
				}
394
				assert.Equal(len(backups), 0)
395
			},
396
		},
397
		{
398
			name:    "open file failed",
399
			baseDir: os.TempDir(),
400
			mock: func(t *testing.T, s Storage, baseDir, modelKey string) {
401
				file, err := os.OpenFile(filepath.Join(baseDir, "download_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
402
				if err != nil {
403
					t.Fatal(err)
404
				}
405
				defer file.Close()
406
				s.(*storage).baseDir = "baz"
407
			},
408
			expect: func(t *testing.T, s Storage, baseDir, modelKey string) {
409
				assert := assert.New(t)
410
				assert.EqualError(s.ClearDownload(modelKey), "remove baz/download_bar.csv: no such file or directory")
411

412
				s.(*storage).baseDir = baseDir
413
				assert.NoError(s.ClearDownload(modelKey))
414
			},
415
		},
416
	}
417

418
	for _, tc := range tests {
419
		t.Run(tc.name, func(t *testing.T) {
420
			s := New(tc.baseDir)
421
			tc.mock(t, s, tc.baseDir, mockModelKey)
422
			tc.expect(t, s, tc.baseDir, mockModelKey)
423
		})
424
	}
425
}
426

427
func TestStorage_ClearNetworkTopology(t *testing.T) {
428
	tests := []struct {
429
		name    string
430
		baseDir string
431
		mock    func(t *testing.T, s Storage, baseDir, modelKey string)
432
		expect  func(t *testing.T, s Storage, baseDir, modelKey string)
433
	}{
434
		{
435
			name:    "clear file",
436
			baseDir: os.TempDir(),
437
			mock: func(t *testing.T, s Storage, baseDir, modelKey string) {
438
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
439
				if err != nil {
440
					t.Fatal(err)
441
				}
442
				defer file.Close()
443
			},
444
			expect: func(t *testing.T, s Storage, baseDir, modelKey string) {
445
				assert := assert.New(t)
446
				assert.NoError(s.ClearNetworkTopology(modelKey))
447
				fileInfos, err := os.ReadDir(filepath.Join(baseDir))
448
				assert.NoError(err)
449

450
				var backups []fs.FileInfo
451
				re := regexp.MustCompile(fmt.Sprintf("%s_%s", NetworkTopologyFilePrefix, modelKey))
452
				for _, fileInfo := range fileInfos {
453
					if !fileInfo.IsDir() && re.MatchString(fileInfo.Name()) {
454
						info, _ := fileInfo.Info()
455
						backups = append(backups, info)
456
					}
457
				}
458
				assert.Equal(len(backups), 0)
459
			},
460
		},
461
		{
462
			name:    "open file failed",
463
			baseDir: os.TempDir(),
464
			mock: func(t *testing.T, s Storage, baseDir, modelKey string) {
465
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
466
				if err != nil {
467
					t.Fatal(err)
468
				}
469
				defer file.Close()
470
				s.(*storage).baseDir = "baz"
471
			},
472
			expect: func(t *testing.T, s Storage, baseDir, modelKey string) {
473
				assert := assert.New(t)
474
				assert.EqualError(s.ClearNetworkTopology(modelKey), "remove baz/networktopology_bar.csv: no such file or directory")
475
				s.(*storage).baseDir = baseDir
476
				assert.NoError(s.ClearNetworkTopology(modelKey))
477
			},
478
		},
479
	}
480

481
	for _, tc := range tests {
482
		t.Run(tc.name, func(t *testing.T) {
483
			s := New(tc.baseDir)
484
			tc.mock(t, s, tc.baseDir, mockModelKey)
485
			tc.expect(t, s, tc.baseDir, mockModelKey)
486
		})
487
	}
488
}
489

490
func TestStorage_Clear(t *testing.T) {
491
	tests := []struct {
492
		name    string
493
		baseDir string
494
		mock    func(t *testing.T, s Storage, baseDir string)
495
		expect  func(t *testing.T, s Storage, baseDir string)
496
	}{
497
		{
498
			name:    "clear file",
499
			baseDir: os.TempDir(),
500
			mock: func(t *testing.T, s Storage, baseDir string) {
501
				s.(*storage).baseDir = filepath.Join(baseDir, "bae")
502
				if err := os.MkdirAll(s.(*storage).baseDir, fs.FileMode(0700)); err != nil {
503
					t.Fatal(err)
504
				}
505

506
				downloadFile, err := os.OpenFile(filepath.Join(s.(*storage).baseDir, "download_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
507
				if err != nil {
508
					t.Fatal(err)
509
				}
510
				defer downloadFile.Close()
511

512
				networkTopologyFile, err := os.OpenFile(filepath.Join(s.(*storage).baseDir, "networktopology_bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
513
				if err != nil {
514
					t.Fatal(err)
515
				}
516
				defer networkTopologyFile.Close()
517
			},
518
			expect: func(t *testing.T, s Storage, baseDir string) {
519
				assert := assert.New(t)
520
				assert.NoError(s.Clear())
521
				_, err := os.Stat(filepath.Join(baseDir, "bae"))
522
				assert.EqualError(err, fmt.Sprintf("stat %s: no such file or directory", filepath.Join(baseDir, "bae")))
523
			},
524
		},
525
	}
526

527
	for _, tc := range tests {
528
		t.Run(tc.name, func(t *testing.T) {
529
			s := New(tc.baseDir)
530
			tc.mock(t, s, tc.baseDir)
531
			tc.expect(t, s, tc.baseDir)
532
		})
533
	}
534
}
535

536
func TestStorage_downloadFilename(t *testing.T) {
537
	baseDir := os.TempDir()
538
	s := New(baseDir)
539

540
	filename := s.(*storage).downloadFilename(mockModelKey)
541
	re := regexp.MustCompile(fmt.Sprintf("%s_%s.%s$", DownloadFilePrefix, mockModelKey, CSVFileExt))
542
	assert := assert.New(t)
543
	assert.True(re.MatchString(filename))
544
}
545

546
func TestStorage_networkTopologyFilename(t *testing.T) {
547
	baseDir := os.TempDir()
548
	s := New(baseDir)
549

550
	filename := s.(*storage).networkTopologyFilename(mockModelKey)
551
	re := regexp.MustCompile(fmt.Sprintf("%s_%s.%s$", NetworkTopologyFilePrefix, mockModelKey, CSVFileExt))
552
	assert := assert.New(t)
553
	assert.True(re.MatchString(filename))
554
}
555