211 lines
6.8 KiB
Go
211 lines
6.8 KiB
Go
package storage
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"io"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/aws/aws-sdk-go-v2/aws"
|
|
"github.com/aws/aws-sdk-go-v2/service/s3"
|
|
"github.com/aws/aws-sdk-go-v2/service/s3/types"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// MockS3Client is a mock implementation of the S3 client interface
|
|
type MockS3Client struct {
|
|
mock.Mock
|
|
}
|
|
|
|
func (m *MockS3Client) PutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
|
|
args := m.Called(ctx, params)
|
|
return args.Get(0).(*s3.PutObjectOutput), args.Error(1)
|
|
}
|
|
|
|
func (m *MockS3Client) GetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error) {
|
|
args := m.Called(ctx, params)
|
|
return args.Get(0).(*s3.GetObjectOutput), args.Error(1)
|
|
}
|
|
|
|
func (m *MockS3Client) DeleteObject(ctx context.Context, params *s3.DeleteObjectInput, optFns ...func(*s3.Options)) (*s3.DeleteObjectOutput, error) {
|
|
args := m.Called(ctx, params)
|
|
return args.Get(0).(*s3.DeleteObjectOutput), args.Error(1)
|
|
}
|
|
|
|
func (m *MockS3Client) ListObjectsV2(ctx context.Context, params *s3.ListObjectsV2Input, optFns ...func(*s3.Options)) (*s3.ListObjectsV2Output, error) {
|
|
args := m.Called(ctx, params)
|
|
return args.Get(0).(*s3.ListObjectsV2Output), args.Error(1)
|
|
}
|
|
|
|
func (m *MockS3Client) HeadObject(ctx context.Context, params *s3.HeadObjectInput, optFns ...func(*s3.Options)) (*s3.HeadObjectOutput, error) {
|
|
args := m.Called(ctx, params)
|
|
return args.Get(0).(*s3.HeadObjectOutput), args.Error(1)
|
|
}
|
|
|
|
func TestS3Storage(t *testing.T) {
|
|
ctx := context.Background()
|
|
mockClient := new(MockS3Client)
|
|
storage := NewS3Storage(mockClient, "test-bucket", "", false)
|
|
|
|
t.Run("Save", func(t *testing.T) {
|
|
mockClient.ExpectedCalls = nil
|
|
mockClient.Calls = nil
|
|
|
|
content := []byte("test content")
|
|
reader := bytes.NewReader(content)
|
|
|
|
// Mock HeadObject to return NotFound error
|
|
mockClient.On("HeadObject", ctx, mock.MatchedBy(func(input *s3.HeadObjectInput) bool {
|
|
return aws.ToString(input.Bucket) == "test-bucket"
|
|
})).Return(&s3.HeadObjectOutput{}, &types.NoSuchKey{
|
|
Message: aws.String("The specified key does not exist."),
|
|
})
|
|
|
|
mockClient.On("PutObject", ctx, mock.MatchedBy(func(input *s3.PutObjectInput) bool {
|
|
return aws.ToString(input.Bucket) == "test-bucket" &&
|
|
aws.ToString(input.ContentType) == "text/plain"
|
|
})).Return(&s3.PutObjectOutput{}, nil)
|
|
|
|
fileInfo, err := storage.Save(ctx, "test.txt", "text/plain", reader)
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, fileInfo.ID)
|
|
assert.Equal(t, "test.txt", fileInfo.Name)
|
|
assert.Equal(t, "text/plain", fileInfo.ContentType)
|
|
|
|
mockClient.AssertExpectations(t)
|
|
})
|
|
|
|
t.Run("Get", func(t *testing.T) {
|
|
content := []byte("test content")
|
|
mockClient.On("GetObject", ctx, mock.MatchedBy(func(input *s3.GetObjectInput) bool {
|
|
return aws.ToString(input.Bucket) == "test-bucket" &&
|
|
aws.ToString(input.Key) == "test-id"
|
|
})).Return(&s3.GetObjectOutput{
|
|
Body: io.NopCloser(bytes.NewReader(content)),
|
|
ContentType: aws.String("text/plain"),
|
|
ContentLength: aws.Int64(int64(len(content))),
|
|
LastModified: aws.Time(time.Now()),
|
|
}, nil)
|
|
|
|
readCloser, info, err := storage.Get(ctx, "test-id")
|
|
require.NoError(t, err)
|
|
defer readCloser.Close()
|
|
|
|
data, err := io.ReadAll(readCloser)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, content, data)
|
|
assert.Equal(t, "test-id", info.ID)
|
|
assert.Equal(t, int64(len(content)), info.Size)
|
|
|
|
mockClient.AssertExpectations(t)
|
|
})
|
|
|
|
t.Run("List", func(t *testing.T) {
|
|
mockClient.ExpectedCalls = nil
|
|
mockClient.Calls = nil
|
|
|
|
mockClient.On("ListObjectsV2", ctx, mock.MatchedBy(func(input *s3.ListObjectsV2Input) bool {
|
|
return aws.ToString(input.Bucket) == "test-bucket" &&
|
|
aws.ToString(input.Prefix) == "test" &&
|
|
aws.ToInt32(input.MaxKeys) == 10
|
|
})).Return(&s3.ListObjectsV2Output{
|
|
Contents: []types.Object{
|
|
{
|
|
Key: aws.String("test1"),
|
|
Size: aws.Int64(100),
|
|
LastModified: aws.Time(time.Now()),
|
|
},
|
|
{
|
|
Key: aws.String("test2"),
|
|
Size: aws.Int64(200),
|
|
LastModified: aws.Time(time.Now()),
|
|
},
|
|
},
|
|
}, nil)
|
|
|
|
// Mock HeadObject for both files
|
|
mockClient.On("HeadObject", ctx, mock.MatchedBy(func(input *s3.HeadObjectInput) bool {
|
|
return aws.ToString(input.Bucket) == "test-bucket" &&
|
|
aws.ToString(input.Key) == "test1"
|
|
})).Return(&s3.HeadObjectOutput{
|
|
ContentType: aws.String("text/plain"),
|
|
Metadata: map[string]string{
|
|
"x-amz-meta-original-name": "test1.txt",
|
|
},
|
|
}, nil).Once()
|
|
|
|
mockClient.On("HeadObject", ctx, mock.MatchedBy(func(input *s3.HeadObjectInput) bool {
|
|
return aws.ToString(input.Bucket) == "test-bucket" &&
|
|
aws.ToString(input.Key) == "test2"
|
|
})).Return(&s3.HeadObjectOutput{
|
|
ContentType: aws.String("text/plain"),
|
|
Metadata: map[string]string{
|
|
"x-amz-meta-original-name": "test2.txt",
|
|
},
|
|
}, nil).Once()
|
|
|
|
files, err := storage.List(ctx, "test", 10, 0)
|
|
require.NoError(t, err)
|
|
assert.Len(t, files, 2)
|
|
assert.Equal(t, "test1", files[0].ID)
|
|
assert.Equal(t, int64(100), files[0].Size)
|
|
assert.Equal(t, "test1.txt", files[0].Name)
|
|
assert.Equal(t, "text/plain", files[0].ContentType)
|
|
|
|
mockClient.AssertExpectations(t)
|
|
})
|
|
|
|
t.Run("Delete", func(t *testing.T) {
|
|
mockClient.On("DeleteObject", ctx, mock.MatchedBy(func(input *s3.DeleteObjectInput) bool {
|
|
return aws.ToString(input.Bucket) == "test-bucket" &&
|
|
aws.ToString(input.Key) == "test-id"
|
|
})).Return(&s3.DeleteObjectOutput{}, nil)
|
|
|
|
err := storage.Delete(ctx, "test-id")
|
|
require.NoError(t, err)
|
|
|
|
mockClient.AssertExpectations(t)
|
|
})
|
|
|
|
t.Run("Exists", func(t *testing.T) {
|
|
mockClient.ExpectedCalls = nil
|
|
mockClient.Calls = nil
|
|
|
|
// Mock HeadObject for existing file
|
|
mockClient.On("HeadObject", ctx, mock.MatchedBy(func(input *s3.HeadObjectInput) bool {
|
|
return aws.ToString(input.Bucket) == "test-bucket" &&
|
|
aws.ToString(input.Key) == "test-id"
|
|
})).Return(&s3.HeadObjectOutput{}, nil).Once()
|
|
|
|
exists, err := storage.Exists(ctx, "test-id")
|
|
require.NoError(t, err)
|
|
assert.True(t, exists)
|
|
|
|
// Mock HeadObject for non-existing file
|
|
mockClient.On("HeadObject", ctx, mock.MatchedBy(func(input *s3.HeadObjectInput) bool {
|
|
return aws.ToString(input.Bucket) == "test-bucket" &&
|
|
aws.ToString(input.Key) == "non-existent"
|
|
})).Return(&s3.HeadObjectOutput{}, &types.NoSuchKey{
|
|
Message: aws.String("The specified key does not exist."),
|
|
}).Once()
|
|
|
|
exists, err = storage.Exists(ctx, "non-existent")
|
|
require.NoError(t, err)
|
|
assert.False(t, exists)
|
|
|
|
mockClient.AssertExpectations(t)
|
|
})
|
|
|
|
t.Run("Custom URL", func(t *testing.T) {
|
|
customStorage := &S3Storage{
|
|
client: mockClient,
|
|
bucket: "test-bucket",
|
|
customURL: "https://custom.domain",
|
|
proxyS3: true,
|
|
}
|
|
assert.Contains(t, customStorage.getObjectURL("test-id"), "https://custom.domain")
|
|
})
|
|
}
|