vaultik/internal/vaultik/remove_snapshot_test.go

package vaultik_test

import (
	"bytes"
	"context"
	"io"
	"strings"
	"sync"
	"testing"

	"github.com/klauspost/compress/zstd"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"sneak.berlin/go/vaultik/internal/log"
	"sneak.berlin/go/vaultik/internal/snapshot"
	"sneak.berlin/go/vaultik/internal/storage"
	"sneak.berlin/go/vaultik/internal/vaultik"
)

// testStorer implements storage.Storer for testing
type testStorer struct {
	mu   sync.Mutex
	data map[string][]byte
}

func newTestStorer() *testStorer {
	return &testStorer{
		data: make(map[string][]byte),
	}
}

func (s *testStorer) Put(_ context.Context, key string, reader io.Reader) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	data, err := io.ReadAll(reader)
	if err != nil {
		return err
	}

	s.data[key] = data

	return nil
}

func (s *testStorer) PutWithProgress(
	ctx context.Context, key string, reader io.Reader,
	_ int64, _ storage.ProgressCallback,
) error {
	return s.Put(ctx, key, reader)
}

func (s *testStorer) Get(_ context.Context, key string) (io.ReadCloser, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	data, exists := s.data[key]
	if !exists {
		return nil, storage.ErrNotFound
	}

	return io.NopCloser(bytes.NewReader(data)), nil
}

func (s *testStorer) Stat(_ context.Context, key string) (*storage.ObjectInfo, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	data, exists := s.data[key]
	if !exists {
		return nil, storage.ErrNotFound
	}

	return &storage.ObjectInfo{
		Key:  key,
		Size: int64(len(data)),
	}, nil
}

func (s *testStorer) Delete(_ context.Context, key string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	delete(s.data, key)

	return nil
}

func (s *testStorer) List(_ context.Context, prefix string) ([]string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	var keys []string

	for key := range s.data {
		if prefix == "" || strings.HasPrefix(key, prefix) {
			keys = append(keys, key)
		}
	}

	return keys, nil
}

func (s *testStorer) ListStream(
	_ context.Context, prefix string,
) <-chan storage.ObjectInfo {
	ch := make(chan storage.ObjectInfo)

	go func() {
		defer close(ch)

		s.mu.Lock()
		defer s.mu.Unlock()

		for key, data := range s.data {
			if prefix == "" || strings.HasPrefix(key, prefix) {
				ch <- storage.ObjectInfo{
					Key:  key,
					Size: int64(len(data)),
				}
			}
		}
	}()

	return ch
}

func (s *testStorer) Info() storage.Info {
	return storage.Info{
		Type:     testLabel,
		Location: "memory",
	}
}

func (s *testStorer) hasKey(key string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()

	_, exists := s.data[key]

	return exists
}

func (s *testStorer) keyCount() int {
	s.mu.Lock()
	defer s.mu.Unlock()

	return len(s.data)
}

// testBlobHashA is a fixture blob hash reused across the remove tests.
const testBlobHashA = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" +
	"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"

// addManifest creates a compressed manifest in storage at the same
// hashed path the production code uses. snapshotID is the human ID;
// the storage path uses RemoteSnapshotKey(id).
func addManifest(
	t *testing.T, store *testStorer, snapshotID string, blobHashes []string,
) {
	t.Helper()

	blobs := make([]snapshot.BlobInfo, len(blobHashes))
	for i, hash := range blobHashes {
		blobs[i] = snapshot.BlobInfo{
			Hash:           hash,
			CompressedSize: 1000,
		}
	}

	remoteKey := snapshot.RemoteSnapshotKey(snapshotID)
	manifest := &snapshot.Manifest{
		SnapshotID: remoteKey,
		BlobCount:  len(blobs),
		Blobs:      blobs,
	}

	data, err := snapshot.EncodeManifest(manifest, 3)
	require.NoError(t, err)

	key := "metadata/" + remoteKey + "/manifest.json.zst"
	err = store.Put(context.Background(), key, bytes.NewReader(data))
	require.NoError(t, err)
}

// remoteKeyPath returns the storage-relative path to a snapshot's
// manifest under the hashed remote-key scheme. Tests use this in
// hasKey/asserts to avoid scattering RemoteSnapshotKey calls
// throughout.
func remoteKeyPath(snapshotID string) string {
	return "metadata/" + snapshot.RemoteSnapshotKey(snapshotID) +
		"/manifest.json.zst"
}

// addBlob adds a fake blob to storage
func addBlob(t *testing.T, store *testStorer, hash string) {
	t.Helper()

	// Create zstd compressed data
	var buf bytes.Buffer

	writer, _ := zstd.NewWriter(&buf)
	_, _ = writer.Write([]byte("blob data"))
	_ = writer.Close()

	key := "blobs/" + hash[:2] + "/" + hash[2:4] + "/" + hash
	err := store.Put(context.Background(), key, bytes.NewReader(buf.Bytes()))
	require.NoError(t, err)
}

// ============================================================================
// Unit Tests for RemoveSnapshot
// ============================================================================

// TestRemoveSnapshot_LocalOnly_PreservesRemote confirms that
// --local-only opts out of the remote-cleanup half: the snapshot is
// removed from the local index, but the remote metadata and blobs are
// untouched.
func TestRemoveSnapshot_LocalOnly_PreservesRemote(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	store := newTestStorer()

	blobA := testBlobHashA
	addManifest(t, store, "snapshot-001", []string{blobA})
	addBlob(t, store, blobA)

	tv := vaultik.NewForTesting(store)

	opts := &vaultik.RemoveOptions{Force: true, LocalOnly: true}
	result, err := tv.RemoveSnapshot("snapshot-001", opts)

	require.NoError(t, err)
	assert.Equal(t, "snapshot-001", result.SnapshotID)
	assert.False(t, result.RemoteRemoved)

	assert.True(t, store.hasKey("blobs/aa/aa/"+blobA))
	assert.True(t, store.hasKey(remoteKeyPath("snapshot-001")))

	assert.Contains(t, tv.Stdout.String(),
		"Removed snapshot 'snapshot-001' from local database")
}

// TestRemoveSnapshot_DefaultRemovesMetadataNotBlobs is the canonical
// case: no flags. The local-DB entry and the snapshot's remote metadata
// are removed; blobs stay on the destination store. The user must then
// run `vaultik prune` to delete blobs no longer referenced by any
// remaining remote manifest, and the output prints that exact command.
func TestRemoveSnapshot_DefaultRemovesMetadataNotBlobs(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	store := newTestStorer()

	blobUnique := testBlobHashA
	blobShared := "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"

	addManifest(t, store, "snapshot-001", []string{blobUnique, blobShared})
	addManifest(t, store, "snapshot-002", []string{blobShared})
	addBlob(t, store, blobUnique)
	addBlob(t, store, blobShared)

	tv := vaultik.NewForTesting(store)

	opts := &vaultik.RemoveOptions{Force: true}
	result, err := tv.RemoveSnapshot("snapshot-001", opts)

	require.NoError(t, err)
	assert.Equal(t, "snapshot-001", result.SnapshotID)
	assert.True(t, result.RemoteRemoved)

	assert.False(t, store.hasKey(remoteKeyPath("snapshot-001")))
	assert.True(t, store.hasKey(remoteKeyPath("snapshot-002")))
	// Blobs are intentionally NOT touched — that's what `vaultik prune`
	// is for.
	assert.True(t, store.hasKey("blobs/aa/aa/"+blobUnique))
	assert.True(t, store.hasKey("blobs/bb/bb/"+blobShared))

	out := tv.Stdout.String()
	assert.Contains(t, out, "Removed snapshot 'snapshot-001' from local database")
	assert.Contains(t, out, "Removed snapshot metadata from remote storage")
	assert.Contains(t, out, "vaultik prune")
}

func TestRemoveSnapshot_DryRun(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	store := newTestStorer()

	blobA := testBlobHashA
	addManifest(t, store, "snapshot-001", []string{blobA})
	addBlob(t, store, blobA)

	initialCount := store.keyCount()

	tv := vaultik.NewForTesting(store)

	opts := &vaultik.RemoveOptions{Force: true, DryRun: true}
	result, err := tv.RemoveSnapshot("snapshot-001", opts)

	require.NoError(t, err)
	assert.True(t, result.DryRun)

	assert.Equal(t, initialCount, store.keyCount())
	assert.True(t, store.hasKey("blobs/aa/aa/"+blobA))
	assert.True(t, store.hasKey(remoteKeyPath("snapshot-001")))

	assert.Contains(t, tv.Stdout.String(), "[Dry run - no changes made]")
}

func TestRemoveAllSnapshots_RequiresForce(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	store := newTestStorer()
	addManifest(t, store, "snapshot-001", []string{})
	addManifest(t, store, "snapshot-002", []string{})

	tv := vaultik.NewForTesting(store)

	opts := &vaultik.RemoveOptions{} // No Force
	_, err := tv.RemoveAllSnapshots(opts)

	require.Error(t, err)
	assert.Contains(t, err.Error(), "--all requires --force")
}

func TestRemoveAllSnapshots_WithForce(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	store := newTestStorer()

	blobA := testBlobHashA
	addManifest(t, store, "snapshot-001", []string{blobA})
	addManifest(t, store, "snapshot-002", []string{blobA})
	addBlob(t, store, blobA)

	tv := vaultik.NewForTesting(store)

	opts := &vaultik.RemoveOptions{Force: true}
	result, err := tv.RemoveAllSnapshots(opts)

	require.NoError(t, err)
	assert.Len(t, result.SnapshotsRemoved, 2)
	assert.True(t, result.RemoteRemoved)

	// Blobs intentionally preserved — that's prune's job.
	assert.True(t, store.hasKey("blobs/aa/aa/"+blobA))
	assert.False(t, store.hasKey(remoteKeyPath("snapshot-001")))
	assert.False(t, store.hasKey(remoteKeyPath("snapshot-002")))

	out := tv.Stdout.String()
	assert.Contains(t, out, "Removed 2 snapshot(s)")
	assert.Contains(t, out, "Removed snapshot metadata from remote storage")
	assert.Contains(t, out, "vaultik prune")
}

func TestRemoveAllSnapshots_DryRun(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	store := newTestStorer()
	addManifest(t, store, "snapshot-001", []string{})
	addManifest(t, store, "snapshot-002", []string{})

	initialCount := store.keyCount()

	tv := vaultik.NewForTesting(store)

	// Default (no LocalOnly) enumerates the orphan remote keys, which
	// matches what NewForTesting has — local DB is empty, so the two
	// addManifest calls land as orphan remote keys.
	opts := &vaultik.RemoveOptions{Force: true, DryRun: true}
	result, err := tv.RemoveAllSnapshots(opts)

	require.NoError(t, err)
	assert.True(t, result.DryRun)
	assert.Len(t, result.SnapshotsRemoved, 2)

	assert.Equal(t, initialCount, store.keyCount())

	assert.Contains(t, tv.Stdout.String(), "[Dry run - no changes made]")
}

func TestRemoveAllSnapshots_NoSnapshots(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	store := newTestStorer()
	// No snapshots added

	tv := vaultik.NewForTesting(store)

	opts := &vaultik.RemoveOptions{Force: true}
	result, err := tv.RemoveAllSnapshots(opts)

	require.NoError(t, err)
	assert.Empty(t, result.SnapshotsRemoved)

	// Verify output
	assert.Contains(t, tv.Stdout.String(), "No snapshots found")
}