tamabox/internal/storage/s3_retry_test.go

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())
}