tamabox/internal/cmd/backup_restore_integration_test.go

package cmd

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/stretchr/testify/require"

	"github.com/tamakyi/TamaBox/internal/conf"
	"github.com/tamakyi/TamaBox/internal/storage"
)

// TestR2BackupRestoreIntegration 默认跳过。只有显式设置 TB_R2_INTEGRATION=1，
// 并通过 NEKOBOX_CONFIG_PATH 提供真实 R2 配置时才运行。
// 测试只访问 tamabox-backup-test/<随机>/ 下的 source/ 与 target/ 两个专属前缀，
// 绝不执行全桶 List，也不会读取、覆盖或删除其他 key。
func TestR2BackupRestoreIntegration(t *testing.T) {
	if os.Getenv("TB_R2_INTEGRATION") != "1" {
		t.Skip("set TB_R2_INTEGRATION=1 to run the real R2 integration test")
	}
	if os.Getenv("NEKOBOX_CONFIG_PATH") == "" {
		t.Fatal("NEKOBOX_CONFIG_PATH must point to the R2 integration configuration")
	}
	require.NoError(t, conf.Init())

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()
	client, err := storage.NewS3Client(
		ctx,
		conf.Upload.ImageEndpoint,
		conf.Upload.ImageBucket,
		conf.Upload.ImageAccessID,
		conf.Upload.ImageAccessSecret,
	)
	require.NoError(t, err)

	random := make([]byte, 12)
	_, err = rand.Read(random)
	require.NoError(t, err)
	basePrefix := fmt.Sprintf("tamabox-backup-test/%d-%s/", time.Now().UTC().UnixNano(), hex.EncodeToString(random))
	sourcePrefix := basePrefix + "source/"
	targetPrefix := basePrefix + "target/"

	objects := []struct {
		relativeKey  string
		body         []byte
		contentType  string
		cacheControl string
	}{
		{
			relativeKey:  "2026/11/04/root.webp",
			body:         []byte("tamabox-r2-integration-root"),
			contentType:  "image/webp",
			cacheControl: "public, max-age=60",
		},
		{
			relativeKey:  "picture/2026/11/04/picture.webp",
			body:         []byte("tamabox-r2-integration-picture"),
			contentType:  "image/webp",
			cacheControl: "public, max-age=120",
		},
	}

	// created 只记录本测试成功创建的精确 key。任何失败清理也只删除这些 key。
	created := make(map[string]bool)
	t.Cleanup(func() {
		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 2*time.Minute)
		defer cleanupCancel()
		for key := range created {
			if err := storage.DeleteObject(cleanupCtx, client, conf.Upload.ImageBucket, key); err != nil {
				t.Errorf("clean up own R2 key %q: %v", key, err)
			}
		}
	})

	for _, object := range objects {
		key := sourcePrefix + object.relativeKey
		err := storage.PutObject(ctx, client, &s3.PutObjectInput{
			Bucket:        aws.String(conf.Upload.ImageBucket),
			Key:           aws.String(key),
			Body:          bytes.NewReader(object.body),
			ContentLength: aws.Int64(int64(len(object.body))),
			ContentType:   aws.String(object.contentType),
			CacheControl:  aws.String(object.cacheControl),
		})
		require.NoError(t, err)
		created[key] = true
	}

	backupDir := t.TempDir()
	require.NoError(t, backupImagesWithClient(ctx, client, conf.Upload.ImageBucket, backupDir, sourcePrefix))
	manifest, err := readImageManifest(backupDir + string(os.PathSeparator) + "images-manifest.json")
	require.NoError(t, err)
	require.Len(t, manifest, len(objects))
	manifestKeys := make(map[string]bool, len(manifest))
	for _, entry := range manifest {
		manifestKeys[entry.Key] = true
	}
	for _, object := range objects {
		require.True(t, manifestKeys[sourcePrefix+object.relativeKey])
	}

	for _, object := range objects {
		sourceKey := sourcePrefix + object.relativeKey
		require.NoError(t, storage.DeleteObject(ctx, client, conf.Upload.ImageBucket, sourceKey))
		delete(created, sourceKey)
	}

	store := &integrationTrackingRestoreStore{
		base: &r2RestoreObjectStore{client: client},
		onUpload: func(key string) {
			created[key] = true
		},
	}
	require.NoError(t, restoreImagesWithStore(ctx, store, conf.Upload.ImageBucket, restoreImagesOptions{
		Input:        backupDir,
		Yes:          true,
		TargetPrefix: targetPrefix,
	}))

	for _, object := range objects {
		targetKey := targetPrefix + sourcePrefix + object.relativeKey
		require.True(t, created[targetKey], "restore must report its own successful Put")
		output, err := client.GetObject(ctx, &s3.GetObjectInput{
			Bucket: aws.String(conf.Upload.ImageBucket),
			Key:    aws.String(targetKey),
		})
		require.NoError(t, err)
		body, readErr := io.ReadAll(output.Body)
		closeErr := output.Body.Close()
		require.NoError(t, readErr)
		require.NoError(t, closeErr)
		digest := sha256.Sum256(object.body)
		require.Equal(t, fmt.Sprintf("%x", digest), fmt.Sprintf("%x", sha256.Sum256(body)))
		require.Equal(t, object.contentType, aws.ToString(output.ContentType))
		require.Equal(t, object.cacheControl, aws.ToString(output.CacheControl))
	}

	for key := range created {
		require.NoError(t, storage.DeleteObject(ctx, client, conf.Upload.ImageBucket, key))
		delete(created, key)
	}
	requireOwnPrefixEmpty(ctx, t, client, conf.Upload.ImageBucket, sourcePrefix)
	requireOwnPrefixEmpty(ctx, t, client, conf.Upload.ImageBucket, targetPrefix)
}

type integrationTrackingRestoreStore struct {
	base     restoreObjectStore
	onUpload func(string)
}

func (s *integrationTrackingRestoreStore) HeadObject(ctx context.Context, input *s3.HeadObjectInput, options ...func(*s3.Options)) (*s3.HeadObjectOutput, error) {
	return s.base.HeadObject(ctx, input, options...)
}

func (s *integrationTrackingRestoreStore) PutObject(ctx context.Context, input *s3.PutObjectInput) error {
	if err := s.base.PutObject(ctx, input); err != nil {
		return err
	}
	s.onUpload(aws.ToString(input.Key))
	return nil
}

func requireOwnPrefixEmpty(ctx context.Context, t *testing.T, client *s3.Client, bucket, prefix string) {
	t.Helper()
	require.NotEmpty(t, prefix)
	output, err := client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
		Bucket: aws.String(bucket),
		Prefix: aws.String(prefix),
	})
	require.NoError(t, err)
	require.Empty(t, output.Contents, "专属 prefix 清理后必须为空：%q", prefix)
	require.False(t, aws.ToBool(output.IsTruncated), "随机专属 prefix 下对象数不应触发分页")
}