vaultik/internal/vaultik/restore_sweeper_test.go

package vaultik //nolint:testpackage // reads the restore blob cache's peak size

import (
	"context"
	"fmt"
	"io"
	"path/filepath"
	"testing"

	"github.com/spf13/afero"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"sneak.berlin/go/vaultik/internal/config"
	"sneak.berlin/go/vaultik/internal/log"
	"sneak.berlin/go/vaultik/internal/storage"
	"sneak.berlin/go/vaultik/internal/ui"
)

// TestRestoreFreesBlobsNoPendingFileNeeds asserts that the restore cache
// frees a blob once no file still to be restored in this run needs it.
// Every original file has a copy with the same content, so every blob
// holds chunks of copies: a restore that leaves the copies out of its
// paths, or fails on them, must not keep the blobs for them.
func TestRestoreFreesBlobsNoPendingFileNeeds(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	// Small files keep the backup and both restores quick.
	const fileBytes = 16 * 1024

	fs := afero.NewOsFs()
	tempDir := t.TempDir()

	dataDir := filepath.Join(tempDir, "source")
	originalsDir := filepath.Join(dataDir, "originals")
	copiesDir := filepath.Join(dataDir, "copies")
	storeDir := t.TempDir()
	dbPath := filepath.Join(tempDir, "index.sqlite")

	originals := writeFilesWithCopies(t, fs, originalsDir, copiesDir, fileBytes)

	ctx := context.Background()

	cfg, storer, snapshotID := setupLocalityBackup(
		ctx, t, fs, dataDir, storeDir, dbPath, fileBytes)

	t.Run("restore by path", func(t *testing.T) {
		t.Parallel()

		restoreDir := filepath.Join(tempDir, "by-path")

		peak, err := restorePeakCachedBlobs(ctx, t, fs, cfg, storer,
			&RestoreOptions{
				SnapshotID: snapshotID,
				TargetDir:  restoreDir,
				Paths:      []string{originalsDir},
			})
		require.NoError(t, err)

		requireRestoredFiles(t, fs, restoreDir, originals)
		assert.LessOrEqualf(t, peak, 1,
			"peak cached blobs was %d; blobs were kept for unselected files",
			peak)
	})

	t.Run("restore with failed files", func(t *testing.T) {
		t.Parallel()

		restoreDir := filepath.Join(tempDir, "with-failures")

		// A regular file where the copies directory belongs makes every
		// copy fail to restore.
		blocker := filepath.Join(restoreDir, copiesDir)
		require.NoError(t, fs.MkdirAll(filepath.Dir(blocker), 0o755))
		require.NoError(t, afero.WriteFile(fs, blocker, nil, 0o644))

		peak, err := restorePeakCachedBlobs(ctx, t, fs, cfg, storer,
			&RestoreOptions{
				SnapshotID: snapshotID,
				TargetDir:  restoreDir,
				SkipErrors: true,
			})
		require.ErrorIs(t, err, errFilesFailedRestore)

		requireRestoredFiles(t, fs, restoreDir, originals)
		assert.LessOrEqualf(t, peak, 1,
			"peak cached blobs was %d; blobs were kept for failed files", peak)
	})
}

// writeFilesWithCopies writes 15 random files of fileBytes into
// originalsDir and a copy of each into copiesDir. With
// setupLocalityBackup given the same size, each file is one chunk and the
// 15 chunks fill 3 blobs. Returns the originals' content by path.
func writeFilesWithCopies(
	t *testing.T, fs afero.Fs, originalsDir, copiesDir string, fileBytes int,
) map[string][]byte {
	t.Helper()

	const fileCount = 15

	require.NoError(t, fs.MkdirAll(originalsDir, 0o755))
	require.NoError(t, fs.MkdirAll(copiesDir, 0o755))

	originals := make(map[string][]byte, fileCount)

	for i := range fileCount {
		name := fmt.Sprintf("file-%02d.bin", i)
		data := randomBytes(t, fileBytes)
		original := filepath.Join(originalsDir, name)
		originals[original] = data

		require.NoError(t, afero.WriteFile(fs, original, data, 0o644))
		require.NoError(t, afero.WriteFile(
			fs, filepath.Join(copiesDir, name), data, 0o644))
	}

	return originals
}

// requireRestoredFiles requires every file in want to have been restored
// under restoreDir with its content intact.
func requireRestoredFiles(
	t *testing.T, fs afero.Fs, restoreDir string, want map[string][]byte,
) {
	t.Helper()

	for path, data := range want {
		got, err := afero.ReadFile(fs, filepath.Join(restoreDir, path))
		require.NoErrorf(t, err, "file missing after restore: %s", path)
		require.Equalf(t, data, got, "byte mismatch for %s", path)
	}
}

// restorePeakCachedBlobs runs a restore and returns the most blobs its
// cache held at once, together with the restore's error.
func restorePeakCachedBlobs(
	ctx context.Context, t *testing.T, fs afero.Fs,
	cfg *config.Config, storer storage.Storer, opts *RestoreOptions,
) (int, error) {
	t.Helper()

	var cache *blobDiskCache

	v := &Vaultik{
		Config:  cfg,
		Storage: storer,
		Fs:      fs,
		Stdout:  io.Discard,
		Stderr:  io.Discard,
		UI:      ui.NewWithColor(io.Discard, false),
		restoreCacheObserver: func(c *blobDiskCache) {
			cache = c
		},
	}
	v.SetContext(ctx)

	err := v.Restore(opts)

	require.NotNil(t, cache, "restoreCacheObserver must fire during restore")

	return cache.PeakLen(), err
}