at-container-registry/pkg/s3/mock.go

package s3

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	awss3 "github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/google/uuid"
)

// MockS3Client implements S3Client for testing without real S3 credentials.
// It generates fake presigned URLs that point to a test server.
type MockS3Client struct {
	// TestServerURL is the base URL for generating fake presigned URLs.
	// Requests to these URLs should be handled by a test server (httptest.Server).
	TestServerURL string

	// UploadID is returned by CreateMultipartUpload.
	// If empty, a UUID is generated.
	UploadID string

	// Track calls for verification in tests
	mu                   sync.Mutex
	CreateMultipartCalls []CreateMultipartCall
	CompleteCalls        []CompleteCall
	AbortCalls           []AbortCall
	UploadPartCalls      []UploadPartCall
	GetObjectCalls       []GetObjectCall
	HeadObjectCalls      []HeadObjectCall
	PutObjectCalls       []PutObjectCall

	// Error injection for testing error handling
	CreateMultipartError error
	CompleteError        error
	AbortError           error
}

// CreateMultipartCall records a CreateMultipartUpload call
type CreateMultipartCall struct {
	Bucket string
	Key    string
}

// CompleteCall records a CompleteMultipartUpload call
type CompleteCall struct {
	Bucket   string
	Key      string
	UploadID string
	Parts    int
}

// AbortCall records an AbortMultipartUpload call
type AbortCall struct {
	Bucket   string
	Key      string
	UploadID string
}

// UploadPartCall records a PresignUploadPart call
type UploadPartCall struct {
	Bucket     string
	Key        string
	UploadID   string
	PartNumber int32
}

// GetObjectCall records a PresignGetObject call
type GetObjectCall struct {
	Bucket string
	Key    string
}

// HeadObjectCall records a PresignHeadObject call
type HeadObjectCall struct {
	Bucket string
	Key    string
}

// PutObjectCall records a PresignPutObject call
type PutObjectCall struct {
	Bucket string
	Key    string
}

// NewMockS3Client creates a new mock S3 client for testing
func NewMockS3Client(testServerURL string) *MockS3Client {
	return &MockS3Client{
		TestServerURL:        testServerURL,
		CreateMultipartCalls: []CreateMultipartCall{},
		CompleteCalls:        []CompleteCall{},
		AbortCalls:           []AbortCall{},
		UploadPartCalls:      []UploadPartCall{},
		GetObjectCalls:       []GetObjectCall{},
		HeadObjectCalls:      []HeadObjectCall{},
		PutObjectCalls:       []PutObjectCall{},
	}
}

// CreateMultipartUpload implements S3Client
func (m *MockS3Client) CreateMultipartUpload(ctx context.Context, input *awss3.CreateMultipartUploadInput, opts ...func(*awss3.Options)) (*awss3.CreateMultipartUploadOutput, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.CreateMultipartCalls = append(m.CreateMultipartCalls, CreateMultipartCall{
		Bucket: aws.ToString(input.Bucket),
		Key:    aws.ToString(input.Key),
	})

	if m.CreateMultipartError != nil {
		return nil, m.CreateMultipartError
	}

	uploadID := m.UploadID
	if uploadID == "" {
		uploadID = "mock-upload-" + uuid.New().String()
	}

	return &awss3.CreateMultipartUploadOutput{
		UploadId: aws.String(uploadID),
	}, nil
}

// CompleteMultipartUpload implements S3Client
func (m *MockS3Client) CompleteMultipartUpload(ctx context.Context, input *awss3.CompleteMultipartUploadInput, opts ...func(*awss3.Options)) (*awss3.CompleteMultipartUploadOutput, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	partsCount := 0
	if input.MultipartUpload != nil {
		partsCount = len(input.MultipartUpload.Parts)
	}

	m.CompleteCalls = append(m.CompleteCalls, CompleteCall{
		Bucket:   aws.ToString(input.Bucket),
		Key:      aws.ToString(input.Key),
		UploadID: aws.ToString(input.UploadId),
		Parts:    partsCount,
	})

	if m.CompleteError != nil {
		return nil, m.CompleteError
	}

	// Return a mock ETag
	etag := "\"mock-etag-" + uuid.New().String() + "\""
	return &awss3.CompleteMultipartUploadOutput{
		ETag: aws.String(etag),
	}, nil
}

// AbortMultipartUpload implements S3Client
func (m *MockS3Client) AbortMultipartUpload(ctx context.Context, input *awss3.AbortMultipartUploadInput, opts ...func(*awss3.Options)) (*awss3.AbortMultipartUploadOutput, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.AbortCalls = append(m.AbortCalls, AbortCall{
		Bucket:   aws.ToString(input.Bucket),
		Key:      aws.ToString(input.Key),
		UploadID: aws.ToString(input.UploadId),
	})

	if m.AbortError != nil {
		return nil, m.AbortError
	}

	return &awss3.AbortMultipartUploadOutput{}, nil
}

// PresignUploadPart implements S3Client
// Returns a mock presigned URL for test server
func (m *MockS3Client) PresignUploadPart(ctx context.Context, input *awss3.UploadPartInput, expires time.Duration) (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.UploadPartCalls = append(m.UploadPartCalls, UploadPartCall{
		Bucket:     aws.ToString(input.Bucket),
		Key:        aws.ToString(input.Key),
		UploadID:   aws.ToString(input.UploadId),
		PartNumber: aws.ToInt32(input.PartNumber),
	})

	url := fmt.Sprintf("%s/upload/%s?partNumber=%d&uploadId=%s",
		m.TestServerURL,
		aws.ToString(input.Key),
		aws.ToInt32(input.PartNumber),
		aws.ToString(input.UploadId))

	return url, nil
}

// PresignGetObject implements S3Client
func (m *MockS3Client) PresignGetObject(ctx context.Context, input *awss3.GetObjectInput, expires time.Duration) (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.GetObjectCalls = append(m.GetObjectCalls, GetObjectCall{
		Bucket: aws.ToString(input.Bucket),
		Key:    aws.ToString(input.Key),
	})

	url := fmt.Sprintf("%s/get/%s", m.TestServerURL, aws.ToString(input.Key))
	return url, nil
}

// PresignHeadObject implements S3Client
func (m *MockS3Client) PresignHeadObject(ctx context.Context, input *awss3.HeadObjectInput, expires time.Duration) (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.HeadObjectCalls = append(m.HeadObjectCalls, HeadObjectCall{
		Bucket: aws.ToString(input.Bucket),
		Key:    aws.ToString(input.Key),
	})

	url := fmt.Sprintf("%s/head/%s", m.TestServerURL, aws.ToString(input.Key))
	return url, nil
}

// PresignPutObject implements S3Client
func (m *MockS3Client) PresignPutObject(ctx context.Context, input *awss3.PutObjectInput, expires time.Duration) (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.PutObjectCalls = append(m.PutObjectCalls, PutObjectCall{
		Bucket: aws.ToString(input.Bucket),
		Key:    aws.ToString(input.Key),
	})

	url := fmt.Sprintf("%s/put/%s", m.TestServerURL, aws.ToString(input.Key))
	return url, nil
}