block uploader restore UT

Signed-off-by: Lyndon-Li <lyonghui@vmware.com>
2026-07-28 19:12:43 +00:00 · 2026-07-24 13:37:40 +08:00
parent 97978ed9b7
commit 236c45b443
1 changed files with 228 additions and 0 deletions
@@ -459,3 +459,231 @@ func TestLoadObjectFromSnapshot(t *testing.T) {
 		})
 	}
 }
+
+func TestGetSourceSize(t *testing.T) {
+	testCases := []struct {
+		name      string
+		snapshot  udmrepo.Snapshot
+		expectErr bool
+		expected  int64
+	}{
+		{
+			name:      "nil tags",
+			snapshot:  udmrepo.Snapshot{},
+			expectErr: true,
+		},
+		{
+			name: "missing tag",
+			snapshot: udmrepo.Snapshot{
+				Tags: map[string]string{},
+			},
+			expectErr: true,
+		},
+		{
+			name: "invalid tag value",
+			snapshot: udmrepo.Snapshot{
+				Tags: map[string]string{
+					"bdev-source-size": "abc",
+				},
+			},
+			expectErr: true,
+		},
+		{
+			name: "valid tag value",
+			snapshot: udmrepo.Snapshot{
+				Tags: map[string]string{
+					"bdev-source-size": "1048576",
+				},
+			},
+			expectErr: false,
+			expected:  1048576,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			size, err := getSourceSize(tc.snapshot)
+			if tc.expectErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tc.expected, size)
+			}
+		})
+	}
+}
+
+func TestFlushZeroBlocks(t *testing.T) {
+	t.Run("success via write fallback", func(t *testing.T) {
+		f, err := os.CreateTemp("", "zerotest-*")
+		require.NoError(t, err)
+		defer os.Remove(f.Name())
+		defer f.Close()
+
+		require.NoError(t, f.Truncate(2048))
+
+		bu := &blockUploader{
+			log: logrus.New(),
+		}
+		bu.log.(*logrus.Logger).Out = io.Discard
+
+		zeroBlock := make([]byte, 1024)
+		err = bu.flushZeroBlocks(f, 0, 2048, zeroBlock, f.Name())
+
+		assert.NoError(t, err)
+
+		data, err := os.ReadFile(f.Name())
+		require.NoError(t, err)
+		assert.Equal(t, make([]byte, 2048), data)
+	})
+}
+
+type errReader struct {
+	err error
+}
+
+func (r *errReader) Read(p []byte) (n int, err error) {
+	return 0, r.err
+}
+
+func (r *errReader) Seek(offset int64, whence int) (int64, error) {
+	return 0, nil
+}
+
+func TestRestoreData(t *testing.T) {
+	t.Run("success", func(t *testing.T) {
+		ctx := context.Background()
+		progress := &mockProgressUpdater{}
+		progress.On("UpdateProgress", mock.Anything).Return()
+		bu := &blockUploader{
+			ctx:      ctx,
+			progress: progress,
+			log:      logrus.New(),
+		}
+
+		f, err := os.CreateTemp("", "restoretest-*")
+		require.NoError(t, err)
+		defer os.Remove(f.Name())
+		defer f.Close()
+
+		data := make([]byte, 1048576)
+		for i := range data {
+			data[i] = 1
+		}
+		reader := bytes.NewReader(data)
+
+		iterMock := cbtmocks.NewIterator(t)
+		iterMock.On("Count").Return(uint64(1))
+		iterMock.On("Next").Return(uint64(0), true).Once()
+		iterMock.On("Next").Return(uint64(0), false)
+
+		written, err := bu.restoreData(reader, f, iterMock, 1048576, f.Name())
+		assert.NoError(t, err)
+		assert.Equal(t, int64(1048576), written)
+
+		f.Seek(0, 0)
+		writtenData, err := io.ReadAll(f)
+		require.NoError(t, err)
+		assert.Equal(t, data, writtenData)
+	})
+
+	t.Run("read err", func(t *testing.T) {
+		ctx := context.Background()
+		bu := &blockUploader{
+			ctx: ctx,
+			log: logrus.New(),
+		}
+
+		f, err := os.CreateTemp("", "restoretest-*")
+		require.NoError(t, err)
+		defer os.Remove(f.Name())
+		defer f.Close()
+
+		reader := &errReader{err: errors.New("read error")}
+
+		iterMock := cbtmocks.NewIterator(t)
+		iterMock.On("Count").Return(uint64(1))
+		iterMock.On("Next").Return(uint64(0), true).Once()
+		iterMock.On("Next").Return(uint64(0), false)
+
+		_, err = bu.restoreData(reader, f, iterMock, 1048576, f.Name())
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "read error")
+	})
+}
+
+func TestBlockUploaderRestore(t *testing.T) {
+	t.Run("missing metadata", func(t *testing.T) {
+		ctx := context.Background()
+		repoWriter := udmrepomocks.NewBackupRepo(t)
+		bu := NewUploader(ctx, repoWriter, nil, logrus.New())
+
+		repoWriter.On("ReadMetadata", mock.Anything, udmrepo.ID("root-id")).Return(nil, errors.New("meta not found"))
+
+		iterMock := cbtmocks.NewIterator(t)
+		_, err := bu.Restore(udmrepo.Snapshot{RootObject: udmrepo.ObjectMetadata{ID: "root-id"}}, destInfo{}, iterMock, nil)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "meta not found")
+	})
+
+	t.Run("success", func(t *testing.T) {
+		ctx := context.Background()
+		repoWriter := udmrepomocks.NewBackupRepo(t)
+		progress := &mockProgressUpdater{}
+		progress.On("UpdateProgress", mock.Anything).Return()
+
+		bu := NewUploader(ctx, repoWriter, progress, logrus.New())
+
+		f, err := os.CreateTemp("", "restoretest-*")
+		require.NoError(t, err)
+		defer os.Remove(f.Name())
+		defer f.Close()
+
+		meta := &udmrepo.Metadata{
+			SubObjects: []udmrepo.ObjectMetadata{
+				{
+					ID:   "data-id",
+					Name: "bdev",
+					Size: 1048576,
+				},
+			},
+		}
+
+		repoWriter.On("ReadMetadata", mock.Anything, udmrepo.ID("root-id")).Return(meta, nil)
+
+		objReader := udmrepomocks.NewObjectReader(t)
+		objReader.On("Read", mock.Anything).Run(func(args mock.Arguments) {
+			p := args.Get(0).([]byte)
+			for i := range p {
+				p[i] = 1
+			}
+		}).Return(1048576, io.EOF).Once()
+		objReader.On("Read", mock.Anything).Return(0, io.EOF)
+		objReader.On("Close").Return(nil)
+
+		repoWriter.On("OpenObject", mock.Anything, udmrepo.ID("data-id")).Return(objReader, nil)
+
+		snap := udmrepo.Snapshot{
+			Description: "test snapshot",
+			RootObject:  udmrepo.ObjectMetadata{ID: "root-id"},
+			Tags: map[string]string{
+				"bdev-source-size": "1048576",
+			},
+		}
+
+		dest := destInfo{
+			dev:  f,
+			size: 2048576,
+			path: f.Name(),
+		}
+
+		iterMock := cbtmocks.NewIterator(t)
+		iterMock.On("Count").Return(uint64(1))
+		iterMock.On("Next").Return(uint64(0), true).Once()
+		iterMock.On("Next").Return(uint64(0), false)
+
+		written, err := bu.Restore(snap, dest, iterMock, nil)
+		assert.NoError(t, err)
+		assert.Equal(t, int64(1048576), written)
+	})
+}