storage/s3_test.go

package storage

import (
	"context"
	"errors"
	"io"
	"strings"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
	"github.com/aws/aws-sdk-go-v2/feature/s3/transfermanager"
	"github.com/aws/aws-sdk-go-v2/service/s3"
)

// Mock implementations for testing

type mockUploader struct {
	uploadFunc func(ctx context.Context, input *transfermanager.UploadObjectInput, opts ...func(*transfermanager.Options)) (*transfermanager.UploadObjectOutput, error)
}

func (m *mockUploader) UploadObject(ctx context.Context, input *transfermanager.UploadObjectInput, opts ...func(*transfermanager.Options)) (*transfermanager.UploadObjectOutput, error) {
	return m.uploadFunc(ctx, input, opts...)
}

type mockDirectUploader struct {
	putObjectFunc func(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error)
}

func (m *mockDirectUploader) PutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
	return m.putObjectFunc(ctx, params, optFns...)
}

type mockPresigner struct {
	presignFunc func(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.PresignOptions)) (*v4.PresignedHTTPRequest, error)
}

func (m *mockPresigner) PresignGetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.PresignOptions)) (*v4.PresignedHTTPRequest, error) {
	return m.presignFunc(ctx, params, optFns...)
}

// Test NewS3 constructor

func TestNewS3(t *testing.T) {
	cfg := aws.Config{
		Region: "us-east-1",
	}
	bucket := "test-bucket"

	s3Instance := NewS3(cfg, bucket)

	if s3Instance == nil {
		t.Fatal("Expected S3 instance, got nil")
	}

	if s3Instance.bucket != bucket {
		t.Errorf("Expected bucket %s, got %s", bucket, s3Instance.bucket)
	}

	if !s3Instance.useDirectUpload {
		t.Error("Expected useDirectUpload to be true for NewS3")
	}

	if s3Instance.directSvc == nil {
		t.Error("Expected directSvc to be set")
	}

	if s3Instance.presigner == nil {
		t.Error("Expected presigner to be set")
	}
}

// Test Store with upload manager pattern

func TestStore_WithUploadManager_Success(t *testing.T) {
	testBucket := "test-bucket"
	testPath := "path/to/file.pdf"
	testContent := "test content"
	testContentType := "application/pdf"
	expectedURL := "https://s3.amazonaws.com/test-bucket/path/to/file.pdf?presigned=true"

	mockUploader := &mockUploader{
		uploadFunc: func(ctx context.Context, input *transfermanager.UploadObjectInput, opts ...func(*transfermanager.Options)) (*transfermanager.UploadObjectOutput, error) {
			// Verify input parameters
			if *input.Bucket != testBucket {
				t.Errorf("Expected bucket %s, got %s", testBucket, *input.Bucket)
			}
			if *input.Key != testPath {
				t.Errorf("Expected key %s, got %s", testPath, *input.Key)
			}

			// Read and verify body
			body, err := io.ReadAll(input.Body)
			if err != nil {
				t.Errorf("Failed to read body: %v", err)
			}
			if string(body) != testContent {
				t.Errorf("Expected content %s, got %s", testContent, string(body))
			}

			return &transfermanager.UploadObjectOutput{
				Key: aws.String(testPath),
			}, nil
		},
	}

	mockPresigner := &mockPresigner{
		presignFunc: func(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.PresignOptions)) (*v4.PresignedHTTPRequest, error) {
			// Verify presign parameters
			if *params.Bucket != testBucket {
				t.Errorf("Expected bucket %s, got %s", testBucket, *params.Bucket)
			}
			if *params.Key != testPath {
				t.Errorf("Expected key %s, got %s", testPath, *params.Key)
			}
			if *params.ResponseContentType != testContentType {
				t.Errorf("Expected content type %s, got %s", testContentType, *params.ResponseContentType)
			}

			return &v4.PresignedHTTPRequest{
				URL: expectedURL,
			}, nil
		},
	}

	s3Instance := &S3{
		bucket:          testBucket,
		svc:             mockUploader,
		presigner:       mockPresigner,
		useDirectUpload: false,
	}

	url, err := s3Instance.Store(testPath, strings.NewReader(testContent), testContentType)
	if err != nil {
		t.Fatalf("Expected no error, got %v", err)
	}

	if url != expectedURL {
		t.Errorf("Expected URL %s, got %s", expectedURL, url)
	}
}

func TestStore_WithUploadManager_UploadError(t *testing.T) {
	testBucket := "test-bucket"
	testPath := "path/to/file.pdf"
	testContent := "test content"
	testContentType := "application/pdf"
	expectedError := errors.New("upload failed")

	mockUploader := &mockUploader{
		uploadFunc: func(ctx context.Context, input *transfermanager.UploadObjectInput, opts ...func(*transfermanager.Options)) (*transfermanager.UploadObjectOutput, error) {
			return nil, expectedError
		},
	}

	mockPresigner := &mockPresigner{
		presignFunc: func(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.PresignOptions)) (*v4.PresignedHTTPRequest, error) {
			t.Error("Presigner should not be called when upload fails")
			return nil, nil
		},
	}

	s3Instance := &S3{
		bucket:          testBucket,
		svc:             mockUploader,
		presigner:       mockPresigner,
		useDirectUpload: false,
	}

	url, err := s3Instance.Store(testPath, strings.NewReader(testContent), testContentType)

	if err == nil {
		t.Fatal("Expected error, got nil")
	}

	if err != expectedError {
		t.Errorf("Expected error %v, got %v", expectedError, err)
	}

	if url != "" {
		t.Errorf("Expected empty URL, got %s", url)
	}
}

func TestStore_WithUploadManager_PresignError(t *testing.T) {
	testBucket := "test-bucket"
	testPath := "path/to/file.pdf"
	testContent := "test content"
	testContentType := "application/pdf"
	expectedError := errors.New("presign failed")

	mockUploader := &mockUploader{
		uploadFunc: func(ctx context.Context, input *transfermanager.UploadObjectInput, opts ...func(*transfermanager.Options)) (*transfermanager.UploadObjectOutput, error) {
			return &transfermanager.UploadObjectOutput{
				Key: aws.String(testPath),
			}, nil
		},
	}

	mockPresigner := &mockPresigner{
		presignFunc: func(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.PresignOptions)) (*v4.PresignedHTTPRequest, error) {
			return nil, expectedError
		},
	}

	s3Instance := &S3{
		bucket:          testBucket,
		svc:             mockUploader,
		presigner:       mockPresigner,
		useDirectUpload: false,
	}

	url, err := s3Instance.Store(testPath, strings.NewReader(testContent), testContentType)

	if err == nil {
		t.Fatal("Expected error, got nil")
	}

	if err != expectedError {
		t.Errorf("Expected error %v, got %v", expectedError, err)
	}

	if url != "" {
		t.Errorf("Expected empty URL, got %s", url)
	}
}

// Test Store with direct upload pattern

func TestStore_WithDirectUpload_Success(t *testing.T) {
	testBucket := "test-bucket"
	testPath := "path/to/file.pdf"
	testContent := "test content"
	testContentType := "application/pdf"
	expectedURL := "https://s3.amazonaws.com/test-bucket/path/to/file.pdf?presigned=true"

	mockDirectUploader := &mockDirectUploader{
		putObjectFunc: func(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
			// Verify input parameters
			if *params.Bucket != testBucket {
				t.Errorf("Expected bucket %s, got %s", testBucket, *params.Bucket)
			}
			if *params.Key != testPath {
				t.Errorf("Expected key %s, got %s", testPath, *params.Key)
			}
			if *params.ContentType != testContentType {
				t.Errorf("Expected content type %s, got %s", testContentType, *params.ContentType)
			}

			// Read and verify body
			body, err := io.ReadAll(params.Body)
			if err != nil {
				t.Errorf("Failed to read body: %v", err)
			}
			if string(body) != testContent {
				t.Errorf("Expected content %s, got %s", testContent, string(body))
			}

			return &s3.PutObjectOutput{}, nil
		},
	}

	mockPresigner := &mockPresigner{
		presignFunc: func(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.PresignOptions)) (*v4.PresignedHTTPRequest, error) {
			// Verify presign parameters
			if *params.Bucket != testBucket {
				t.Errorf("Expected bucket %s, got %s", testBucket, *params.Bucket)
			}
			if *params.Key != testPath {
				t.Errorf("Expected key %s, got %s", testPath, *params.Key)
			}

			// Verify 15 minute expiry was requested
			// Note: We can't directly verify the duration in the options, but we can ensure it's called
			return &v4.PresignedHTTPRequest{
				URL: expectedURL,
			}, nil
		},
	}

	s3Instance := &S3{
		bucket:          testBucket,
		directSvc:       mockDirectUploader,
		presigner:       mockPresigner,
		useDirectUpload: true,
	}

	url, err := s3Instance.Store(testPath, strings.NewReader(testContent), testContentType)
	if err != nil {
		t.Fatalf("Expected no error, got %v", err)
	}

	if url != expectedURL {
		t.Errorf("Expected URL %s, got %s", expectedURL, url)
	}
}

func TestStore_WithDirectUpload_PutObjectError(t *testing.T) {
	testBucket := "test-bucket"
	testPath := "path/to/file.pdf"
	testContent := "test content"
	testContentType := "application/pdf"
	expectedError := errors.New("put object failed")

	mockDirectUploader := &mockDirectUploader{
		putObjectFunc: func(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
			return nil, expectedError
		},
	}

	mockPresigner := &mockPresigner{
		presignFunc: func(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.PresignOptions)) (*v4.PresignedHTTPRequest, error) {
			t.Error("Presigner should not be called when PutObject fails")
			return nil, nil
		},
	}

	s3Instance := &S3{
		bucket:          testBucket,
		directSvc:       mockDirectUploader,
		presigner:       mockPresigner,
		useDirectUpload: true,
	}

	url, err := s3Instance.Store(testPath, strings.NewReader(testContent), testContentType)

	if err == nil {
		t.Fatal("Expected error, got nil")
	}

	if err != expectedError {
		t.Errorf("Expected error %v, got %v", expectedError, err)
	}

	if url != "" {
		t.Errorf("Expected empty URL, got %s", url)
	}
}

func TestStore_WithDirectUpload_PresignError(t *testing.T) {
	testBucket := "test-bucket"
	testPath := "path/to/file.pdf"
	testContent := "test content"
	testContentType := "application/pdf"
	expectedError := errors.New("presign failed")

	mockDirectUploader := &mockDirectUploader{
		putObjectFunc: func(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
			return &s3.PutObjectOutput{}, nil
		},
	}

	mockPresigner := &mockPresigner{
		presignFunc: func(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.PresignOptions)) (*v4.PresignedHTTPRequest, error) {
			return nil, expectedError
		},
	}

	s3Instance := &S3{
		bucket:          testBucket,
		directSvc:       mockDirectUploader,
		presigner:       mockPresigner,
		useDirectUpload: true,
	}

	url, err := s3Instance.Store(testPath, strings.NewReader(testContent), testContentType)

	if err == nil {
		t.Fatal("Expected error, got nil")
	}

	if err != expectedError {
		t.Errorf("Expected error %v, got %v", expectedError, err)
	}

	if url != "" {
		t.Errorf("Expected empty URL, got %s", url)
	}
}

// Test that presign expiry is set correctly

func TestStore_PresignExpiry(t *testing.T) {
	testBucket := "test-bucket"
	testPath := "path/to/file.pdf"
	testContent := "test content"
	testContentType := "application/pdf"

	var capturedExpiry time.Duration

	mockUploader := &mockUploader{
		uploadFunc: func(ctx context.Context, input *transfermanager.UploadObjectInput, opts ...func(*transfermanager.Options)) (*transfermanager.UploadObjectOutput, error) {
			return &transfermanager.UploadObjectOutput{Key: aws.String(testPath)}, nil
		},
	}

	mockPresigner := &mockPresigner{
		presignFunc: func(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.PresignOptions)) (*v4.PresignedHTTPRequest, error) {
			// Apply options to capture the expiry
			opts := &s3.PresignOptions{}
			for _, fn := range optFns {
				fn(opts)
			}
			capturedExpiry = opts.Expires

			return &v4.PresignedHTTPRequest{
				URL: "https://example.com/presigned",
			}, nil
		},
	}

	s3Instance := &S3{
		bucket:          testBucket,
		svc:             mockUploader,
		presigner:       mockPresigner,
		useDirectUpload: false,
	}

	_, err := s3Instance.Store(testPath, strings.NewReader(testContent), testContentType)
	if err != nil {
		t.Fatalf("Expected no error, got %v", err)
	}

	expectedExpiry := 15 * time.Minute
	if capturedExpiry != expectedExpiry {
		t.Errorf("Expected presign expiry of %v, got %v", expectedExpiry, capturedExpiry)
	}
}