Add AWS storage unit tests

Signed-off-by: Daniel Weiße <dw@edgeless.systems>

kms/storage/awss3storage.go

@@ -13,10 +13,16 @@ import (
 	"github.com/edgelesssys/constellation/kms/config"
 )
 
+type awsS3ClientAPI interface {
+	GetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error)
+	PutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error)
+	CreateBucket(ctx context.Context, params *s3.CreateBucketInput, optFns ...func(*s3.Options)) (*s3.CreateBucketOutput, error)
+}
+
 // AWSS3Storage is an implementation of the Storage interface, storing keys in AWS S3 buckets.
 type AWSS3Storage struct {
 	bucketID string
-	client   *s3.Client
+	client   awsS3ClientAPI
 	optFns   []func(*s3.Options)
 }
 
@@ -31,22 +37,13 @@ func NewAWSS3Storage(ctx context.Context, bucketID string, optFns ...func(*s3.Op
 	}
 	client := s3.NewFromConfig(cfg, optFns...)
 
+	store := &AWSS3Storage{client: client, bucketID: bucketID, optFns: optFns}
+
 	// Try to create new bucket, continue if bucket already exists
-	createBucketInput := &s3.CreateBucketInput{
+	if err := store.createBucket(ctx, bucketID, cfg.Region, optFns...); err != nil {
-		Bucket: &bucketID,
+		return nil, err
-		CreateBucketConfiguration: &types.CreateBucketConfiguration{
-			LocationConstraint: types.BucketLocationConstraint(cfg.Region),
-		},
 	}
-	_, err = client.CreateBucket(ctx, createBucketInput, optFns...)
+	return store, nil
-	if err != nil {
-		var bne *types.BucketAlreadyExists
-		var baowby *types.BucketAlreadyOwnedByYou
-		if !(errors.As(err, &bne) || errors.As(err, &baowby)) {
-			return nil, fmt.Errorf("creating storage container: %w", err)
-		}
-	}
-	return &AWSS3Storage{client: client, bucketID: bucketID, optFns: optFns}, nil
 }
 
 // Get returns a DEK from from AWS S3 Storage by key ID.
@@ -79,3 +76,20 @@ func (s *AWSS3Storage) Put(ctx context.Context, keyID string, data []byte) error
 	}
 	return nil
 }
+
+func (s *AWSS3Storage) createBucket(ctx context.Context, bucketID, region string, optFns ...func(*s3.Options)) error {
+	createBucketInput := &s3.CreateBucketInput{
+		Bucket: &bucketID,
+		CreateBucketConfiguration: &types.CreateBucketConfiguration{
+			LocationConstraint: types.BucketLocationConstraint(region),
+		},
+	}
+	if _, err := s.client.CreateBucket(ctx, createBucketInput, optFns...); err != nil {
+		var bne *types.BucketAlreadyExists
+		var baowby *types.BucketAlreadyOwnedByYou
+		if !(errors.As(err, &bne) || errors.As(err, &baowby)) {
+			return fmt.Errorf("creating storage container: %w", err)
+		}
+	}
+	return nil
+}

kms/storage/awss3storage_test.go

@@ -0,0 +1,164 @@
+package storage
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"io"
+	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/service/s3"
+	"github.com/aws/aws-sdk-go-v2/service/s3/types"
+	"github.com/stretchr/testify/assert"
+)
+
+type stubAWSS3StorageClient struct {
+	getObjectOutputData []byte
+	getObjectErr        error
+	putObjectErr        error
+	savedObject         []byte
+	createBucketCalled  bool
+	createBucketErr     error
+}
+
+func (s *stubAWSS3StorageClient) GetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error) {
+	return &s3.GetObjectOutput{
+		Body: io.NopCloser(bytes.NewReader(s.getObjectOutputData)),
+	}, s.getObjectErr
+}
+
+func (s *stubAWSS3StorageClient) PutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
+	out, err := io.ReadAll(params.Body)
+	if err != nil {
+		panic(err)
+	}
+	s.savedObject = out
+	return &s3.PutObjectOutput{}, s.putObjectErr
+}
+
+func (s *stubAWSS3StorageClient) CreateBucket(ctx context.Context, params *s3.CreateBucketInput, optFns ...func(*s3.Options)) (*s3.CreateBucketOutput, error) {
+	s.createBucketCalled = true
+	return &s3.CreateBucketOutput{}, s.createBucketErr
+}
+
+func TestAWSS3Get(t *testing.T) {
+	testCases := map[string]struct {
+		client      *stubAWSS3StorageClient
+		unsetError  bool
+		errExpected bool
+	}{
+		"Get successful": {
+			client: &stubAWSS3StorageClient{getObjectOutputData: []byte("test-data")},
+		},
+		"GetObject fails": {
+			client:      &stubAWSS3StorageClient{getObjectErr: errors.New("error")},
+			errExpected: true,
+		},
+		"GetObject fails with NoSuchKey": {
+			client:      &stubAWSS3StorageClient{getObjectErr: &types.NoSuchKey{}},
+			errExpected: true,
+			unsetError:  true,
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+
+			store := &AWSS3Storage{
+				client: tc.client,
+			}
+
+			out, err := store.Get(context.Background(), "test-key")
+			if tc.errExpected {
+				assert.Error(err)
+
+				if tc.unsetError {
+					assert.ErrorIs(err, ErrDEKUnset)
+				} else {
+					assert.False(errors.Is(err, ErrDEKUnset))
+				}
+
+			} else {
+				assert.NoError(err)
+				assert.Equal(tc.client.getObjectOutputData, out)
+			}
+		})
+	}
+}
+
+func TestAWSS3Put(t *testing.T) {
+	testCases := map[string]struct {
+		client      *stubAWSS3StorageClient
+		errExpected bool
+	}{
+		"Put successful": {
+			client: &stubAWSS3StorageClient{},
+		},
+		"PutObject fails": {
+			client:      &stubAWSS3StorageClient{putObjectErr: errors.New("error")},
+			errExpected: true,
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+
+			store := &AWSS3Storage{
+				client: tc.client,
+			}
+
+			testData := []byte{0x1, 0x2, 0x3}
+
+			err := store.Put(context.Background(), "test-key", testData)
+			if tc.errExpected {
+				assert.Error(err)
+			} else {
+				assert.NoError(err)
+				assert.Equal(testData, tc.client.savedObject)
+			}
+		})
+	}
+}
+
+func TestAWSS3CreateBucket(t *testing.T) {
+	testCases := map[string]struct {
+		client      *stubAWSS3StorageClient
+		errExpected bool
+	}{
+		"CreateBucket successful": {
+			client: &stubAWSS3StorageClient{},
+		},
+		"CreateBucket fails": {
+			client:      &stubAWSS3StorageClient{createBucketErr: errors.New("error")},
+			errExpected: true,
+		},
+		"CreateBucket fails with BucketAlreadyExists": {
+			client:      &stubAWSS3StorageClient{createBucketErr: &types.BucketAlreadyExists{}},
+			errExpected: false,
+		},
+		"CreateBucket fails with BucketAlreadyOwnedByYou": {
+			client:      &stubAWSS3StorageClient{createBucketErr: &types.BucketAlreadyOwnedByYou{}},
+			errExpected: false,
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+
+			store := &AWSS3Storage{
+				client: tc.client,
+			}
+
+			err := store.createBucket(context.Background(), "test-bucket", "test-region")
+			if tc.errExpected {
+				assert.Error(err)
+			} else {
+				assert.NoError(err)
+				assert.True(tc.client.createBucketCalled)
+			}
+		})
+	}
+}