package storage import ( "bytes" "context" "io" "net" "net/http" "net/http/httptest" "sync/atomic" "testing" "time" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestPutObjectRetriesAndRewindsBody(t *testing.T) { var attempts atomic.Int32 bodies := make(chan string, 3) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, _ := io.ReadAll(r.Body) bodies <- string(body) if attempts.Add(1) < 3 { http.Error(w, "temporary", http.StatusInternalServerError) return } w.WriteHeader(http.StatusOK) })) defer server.Close() client, err := NewS3Client(context.Background(), server.URL, "images", "id", "secret") require.NoError(t, err) err = PutObject(context.Background(), client, &s3.PutObjectInput{ Bucket: aws.String("images"), Key: aws.String("retry.txt"), Body: bytes.NewReader([]byte("complete-body")), }) require.NoError(t, err) assert.Equal(t, int32(3), attempts.Load()) close(bodies) for body := range bodies { assert.Equal(t, "complete-body", body) } } func TestPutObjectDoesNotRetryPermanentErrors(t *testing.T) { for _, statusCode := range []int{http.StatusForbidden, http.StatusNotFound} { t.Run(http.StatusText(statusCode), func(t *testing.T) { var attempts atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { attempts.Add(1) http.Error(w, "permanent", statusCode) })) defer server.Close() client, err := NewS3Client(context.Background(), server.URL, "images", "id", "secret") require.NoError(t, err) err = PutObject(context.Background(), client, &s3.PutObjectInput{ Bucket: aws.String("images"), Key: aws.String("permanent.txt"), Body: bytes.NewReader([]byte("body")), }) require.Error(t, err) assert.Equal(t, int32(1), attempts.Load()) }) } } func TestPutObjectRetriesTemporaryStatusCodes(t *testing.T) { for _, statusCode := range []int{http.StatusRequestTimeout, http.StatusTooManyRequests, http.StatusBadGateway} { t.Run(http.StatusText(statusCode), func(t *testing.T) { var attempts atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if attempts.Add(1) == 1 { http.Error(w, "temporary", statusCode) return } w.WriteHeader(http.StatusOK) })) defer server.Close() client, err := NewS3Client(context.Background(), server.URL, "images", "id", "secret") require.NoError(t, err) err = PutObject(context.Background(), client, &s3.PutObjectInput{ Bucket: aws.String("images"), Key: aws.String("temporary.txt"), Body: bytes.NewReader([]byte("body")), }) require.NoError(t, err) assert.Equal(t, int32(2), attempts.Load()) }) } } func TestPutObjectStopsAfterThreeAttempts(t *testing.T) { var attempts atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { attempts.Add(1) http.Error(w, "still unavailable", http.StatusServiceUnavailable) })) defer server.Close() client, err := NewS3Client(context.Background(), server.URL, "images", "id", "secret") require.NoError(t, err) err = PutObject(context.Background(), client, &s3.PutObjectInput{ Bucket: aws.String("images"), Key: aws.String("unavailable.txt"), Body: bytes.NewReader([]byte("body")), }) require.Error(t, err) assert.Equal(t, int32(3), attempts.Load()) } func TestPutObjectHonorsCanceledContext(t *testing.T) { var attempts atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { attempts.Add(1) w.WriteHeader(http.StatusOK) })) defer server.Close() client, err := NewS3Client(context.Background(), server.URL, "images", "id", "secret") require.NoError(t, err) ctx, cancel := context.WithCancel(context.Background()) cancel() err = PutObject(ctx, client, &s3.PutObjectInput{ Bucket: aws.String("images"), Key: aws.String("canceled.txt"), Body: bytes.NewReader([]byte("body")), }) require.ErrorIs(t, err, context.Canceled) assert.Zero(t, attempts.Load()) } func TestRetryBackoffCanBeCanceled(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) done := make(chan error, 1) go func() { done <- waitForS3Retry(ctx, time.Hour) }() cancel() select { case err := <-done: require.ErrorIs(t, err, context.Canceled) case <-time.After(time.Second): t.Fatal("退避未响应 context 取消") } } func TestNetworkErrorsAreRetryable(t *testing.T) { assert.True(t, retryableS3Error(&net.DNSError{Err: "temporary resolver failure", IsTimeout: true})) } func TestDeleteObjectNotFoundIsIdempotent(t *testing.T) { var attempts atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { attempts.Add(1) assert.Equal(t, http.MethodDelete, r.Method) http.Error(w, "not found", http.StatusNotFound) })) defer server.Close() client, err := NewS3Client(context.Background(), server.URL, "images", "id", "secret") require.NoError(t, err) require.NoError(t, DeleteObject(context.Background(), client, "images", "already-gone.txt")) assert.Equal(t, int32(1), attempts.Load()) }