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 下对象数不应触发分页") }