Add GCP storage unit tests

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

kms/storage/gcloudstorage.go

@@ -9,8 +9,37 @@ import (
 	"google.golang.org/api/option"
 )
 
+type gcpStorageAPI interface {
+	Attrs(ctx context.Context, bucketName string) (*storage.BucketAttrs, error)
+	Close() error
+	CreateBucket(ctx context.Context, bucketName, projectID string, attrs *storage.BucketAttrs) error
+	NewWriter(ctx context.Context, bucketName, objectName string) io.WriteCloser
+	NewReader(ctx context.Context, bucketName, objectName string) (io.ReadCloser, error)
+}
+
+type wrappedGCPClient struct {
+	*storage.Client
+}
+
+func (c *wrappedGCPClient) Attrs(ctx context.Context, bucketName string) (*storage.BucketAttrs, error) {
+	return c.Client.Bucket(bucketName).Attrs(ctx)
+}
+
+func (c *wrappedGCPClient) CreateBucket(ctx context.Context, bucketName, projectID string, attrs *storage.BucketAttrs) error {
+	return c.Client.Bucket(bucketName).Create(ctx, projectID, attrs)
+}
+
+func (c *wrappedGCPClient) NewWriter(ctx context.Context, bucketName, objectName string) io.WriteCloser {
+	return c.Client.Bucket(bucketName).Object(objectName).NewWriter(ctx)
+}
+
+func (c *wrappedGCPClient) NewReader(ctx context.Context, bucketName, objectName string) (io.ReadCloser, error) {
+	return c.Client.Bucket(bucketName).Object(objectName).NewReader(ctx)
+}
+
 // GoogleCloudStorage is an implementation of the Storage interface, storing keys in Google Cloud Storage buckets.
 type GoogleCloudStorage struct {
+	newClient  func(ctx context.Context, opts ...option.ClientOption) (gcpStorageAPI, error)
 	projectID  string
 	bucketName string
 	opts       []option.ClientOption
@@ -20,40 +49,30 @@ type GoogleCloudStorage struct {
 //
 // The parameter bucketOptions is optional, if not present default options will be created.
 func NewGoogleCloudStorage(ctx context.Context, projectID, bucketName string, bucketOptions *storage.BucketAttrs, opts ...option.ClientOption) (*GoogleCloudStorage, error) {
-	gcStorage := &GoogleCloudStorage{
+	s := &GoogleCloudStorage{
+		newClient:  gcpStorageClientFactory,
 		projectID:  projectID,
 		bucketName: bucketName,
 		opts:       opts,
 	}
 
 	// Make sure the storage bucket exists, if not create it
-	client, err := storage.NewClient(ctx, gcStorage.opts...)
+	if err := s.createContainerOrContinue(ctx, bucketOptions); err != nil {
-	if err != nil {
 		return nil, err
 	}
-	defer client.Close()
 
-	_, err = client.Bucket(gcStorage.bucketName).Attrs(ctx)
+	return s, nil
-	if err == nil {
-		return gcStorage, nil
-	}
-
-	if errors.Is(err, storage.ErrBucketNotExist) {
-		err = client.Bucket(gcStorage.bucketName).Create(ctx, gcStorage.projectID, bucketOptions)
-	}
-
-	return gcStorage, err
 }
 
 // Get returns a DEK from Google Cloud Storage by key ID.
 func (s *GoogleCloudStorage) Get(ctx context.Context, keyID string) ([]byte, error) {
-	client, err := storage.NewClient(ctx, s.opts...)
+	client, err := s.newClient(ctx, s.opts...)
 	if err != nil {
 		return nil, err
 	}
 	defer client.Close()
 
-	reader, err := client.Bucket(s.bucketName).Object(keyID).NewReader(ctx)
+	reader, err := client.NewReader(ctx, s.bucketName, keyID)
 	if err != nil {
 		if errors.Is(err, storage.ErrObjectNotExist) {
 			return nil, ErrDEKUnset
@@ -67,17 +86,39 @@ func (s *GoogleCloudStorage) Get(ctx context.Context, keyID string) ([]byte, err
 
 // Put saves a DEK to Google Cloud Storage by key ID.
 func (s *GoogleCloudStorage) Put(ctx context.Context, keyID string, data []byte) error {
-	client, err := storage.NewClient(ctx, s.opts...)
+	client, err := s.newClient(ctx, s.opts...)
 	if err != nil {
 		return err
 	}
 	defer client.Close()
 
-	writer := client.Bucket(s.bucketName).Object(keyID).NewWriter(ctx)
+	writer := client.NewWriter(ctx, s.bucketName, keyID)
+	defer writer.Close()
 
-	if _, err := writer.Write(data); err != nil {
+	_, err = writer.Write(data)
+	return err
+}
+
+func (s *GoogleCloudStorage) createContainerOrContinue(ctx context.Context, bucketOptions *storage.BucketAttrs) error {
+	client, err := s.newClient(ctx, s.opts...)
+	if err != nil {
+		return err
+	}
+	defer client.Close()
+
+	if _, err := client.Attrs(ctx, s.bucketName); errors.Is(err, storage.ErrBucketNotExist) {
+		return client.CreateBucket(ctx, s.bucketName, s.projectID, bucketOptions)
+	} else if err != nil {
 		return err
 	}
 
-	return writer.Close()
+	return nil
+}
+
+func gcpStorageClientFactory(ctx context.Context, opts ...option.ClientOption) (gcpStorageAPI, error) {
+	client, err := storage.NewClient(ctx, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return &wrappedGCPClient{client}, nil
 }

kms/storage/gcloudstorage_integration_test.go

@@ -0,0 +1,109 @@
+//go:build integration
+
+package storage
+
+import (
+	"context"
+	"io"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/docker/docker/api/types"
+	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/client"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/api/option"
+)
+
+const storageEmulator = "gcr.io/cloud-devrel-public-resources/storage-testbench"
+
+func TestGoogleCloudStorage(t *testing.T) {
+	assert := assert.New(t)
+	require := require.New(t)
+
+	ctx := context.Background()
+
+	// Set up the Storage Emulator
+	t.Log("Creating storage emulator...")
+	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
+	require.NoError(err)
+	emulator, err := setupEmulator(ctx, cli, storageEmulator)
+	require.NoError(err)
+	defer func() { _ = cli.ContainerStop(ctx, emulator.ID, nil) }()
+
+	// Run the actual test
+	t.Setenv("STORAGE_EMULATOR_HOST", "localhost:9000")
+
+	bucketName := "test-bucket"
+	projectName := "test-project"
+
+	t.Log("Running test...")
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*50)
+	defer cancel()
+	storage, err := NewGoogleCloudStorage(ctx, projectName, bucketName, nil, option.WithoutAuthentication())
+	require.NoError(err)
+
+	testDEK1 := []byte("test DEK")
+	testDEK2 := []byte("more test DEK")
+
+	// request unset value
+	_, err = storage.Get(ctx, "test:input")
+	assert.Error(err)
+
+	// test Put method
+	assert.NoError(storage.Put(ctx, "volume01", testDEK1))
+	assert.NoError(storage.Put(ctx, "volume02", testDEK2))
+
+	// make sure values have been set
+	val, err := storage.Get(ctx, "volume01")
+	assert.NoError(err)
+	assert.Equal(testDEK1, val)
+	val, err = storage.Get(ctx, "volume02")
+	assert.NoError(err)
+	assert.Equal(testDEK2, val)
+
+	_, err = storage.Get(ctx, "invalid:key")
+	assert.Error(err)
+	assert.ErrorIs(err, ErrDEKUnset)
+}
+
+func setupEmulator(ctx context.Context, cli *client.Client, imageName string) (container.ContainerCreateCreatedBody, error) {
+	reader, err := cli.ImagePull(ctx, imageName, types.ImagePullOptions{})
+	if err != nil {
+		return container.ContainerCreateCreatedBody{}, err
+	}
+	if _, err := io.Copy(os.Stdout, reader); err != nil {
+		return container.ContainerCreateCreatedBody{}, err
+	}
+	if err := reader.Close(); err != nil {
+		return container.ContainerCreateCreatedBody{}, err
+	}
+
+	// the 3 true statements are necessary to attach later to the container log
+	containerConfig := &container.Config{
+		Image:        storageEmulator,
+		AttachStdout: true,
+		AttachStderr: true,
+		Tty:          true,
+	}
+	emulator, err := cli.ContainerCreate(ctx, containerConfig, &container.HostConfig{NetworkMode: container.NetworkMode("host"), AutoRemove: true}, nil, nil, "google-cloud-storage-test")
+	if err != nil {
+		return emulator, err
+	}
+	err = cli.ContainerStart(ctx, emulator.ID, types.ContainerStartOptions{})
+	if err != nil {
+		return emulator, err
+	}
+
+	logs, err := cli.ContainerLogs(ctx, emulator.ID, types.ContainerLogsOptions{
+		ShowStdout: true,
+		Follow:     true,
+	})
+	if err != nil {
+		return emulator, err
+	}
+	go func() { _, _ = io.Copy(os.Stdout, logs) }()
+	return emulator, nil
+}

kms/storage/gcloudstorage_test.go

@@ -1,107 +1,222 @@
 package storage
 
 import (
+	"bytes"
 	"context"
+	"errors"
 	"io"
-	"os"
 	"testing"
-	"time"
 
-	"github.com/docker/docker/api/types"
+	"cloud.google.com/go/storage"
-	"github.com/docker/docker/api/types/container"
-	"github.com/docker/docker/client"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 	"google.golang.org/api/option"
 )
 
-const storageEmulator = "gcr.io/cloud-devrel-public-resources/storage-testbench"
+type stubGCPStorageAPI struct {
+	newClientErr       error
+	attrsErr           error
+	createBucketErr    error
+	createBucketCalled bool
+	newReaderErr       error
+	newReaderOutput    []byte
+	writer             *stubWriteCloser
+}
 
-func TestGoogleCloudStorage(t *testing.T) {
+func (s *stubGCPStorageAPI) stubClientFactory(ctx context.Context, opts ...option.ClientOption) (gcpStorageAPI, error) {
+	return s, s.newClientErr
+}
+
+func (s *stubGCPStorageAPI) Attrs(ctx context.Context, bucketName string) (*storage.BucketAttrs, error) {
+	return &storage.BucketAttrs{}, s.attrsErr
+}
+
+func (s *stubGCPStorageAPI) Close() error {
+	return nil
+}
+
+func (s *stubGCPStorageAPI) CreateBucket(ctx context.Context, bucketName, projectID string, attrs *storage.BucketAttrs) error {
+	s.createBucketCalled = true
+	return s.createBucketErr
+}
+
+func (s *stubGCPStorageAPI) NewWriter(ctx context.Context, bucketName, objectName string) io.WriteCloser {
+	return s.writer
+}
+
+func (s *stubGCPStorageAPI) NewReader(ctx context.Context, bucketName, objectName string) (io.ReadCloser, error) {
+	return io.NopCloser(bytes.NewReader(s.newReaderOutput)), s.newReaderErr
+}
+
+type stubWriteCloser struct {
+	result   *[]byte
+	writeErr error
+	writeN   int
+}
+
+func (s stubWriteCloser) Write(p []byte) (int, error) {
+	*s.result = p
+	return s.writeN, s.writeErr
+}
+
+func (s stubWriteCloser) Close() error {
+	return nil
+}
+
+func TestGCPGet(t *testing.T) {
+	someErr := errors.New("error")
+
+	testCases := map[string]struct {
+		client      *stubGCPStorageAPI
+		unsetError  bool
+		errExpected bool
+	}{
+		"success": {
+			client: &stubGCPStorageAPI{newReaderOutput: []byte("test-data")},
+		},
+		"creating client fails": {
+			client:      &stubGCPStorageAPI{newClientErr: someErr},
+			errExpected: true,
+		},
+		"NewReader fails": {
+			client:      &stubGCPStorageAPI{newReaderErr: someErr},
+			errExpected: true,
+		},
+		"ErrObjectNotExist error": {
+			client:      &stubGCPStorageAPI{newReaderErr: storage.ErrObjectNotExist},
+			unsetError:  true,
+			errExpected: true,
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
 			assert := assert.New(t)
-	require := require.New(t)
 
-	ctx := context.Background()
+			client := &GoogleCloudStorage{
+				newClient:  tc.client.stubClientFactory,
+				projectID:  "test",
+				bucketName: "test",
+			}
 
-	// Set up the Storage Emulator
+			out, err := client.Get(context.Background(), "test-key")
-	t.Log("Creating storage emulator...")
+			if tc.errExpected {
-	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
-	require.NoError(err)
-	emulator, err := setupEmulator(ctx, cli, storageEmulator)
-	require.NoError(err)
-	defer cli.ContainerStop(ctx, emulator.ID, nil)
-
-	// Run the actual test
-	t.Setenv("STORAGE_EMULATOR_HOST", "localhost:9000")
-
-	bucketName := "test-bucket"
-	projectName := "test-project"
-
-	t.Log("Running test...")
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*50)
-	defer cancel()
-	storage, err := NewGoogleCloudStorage(ctx, projectName, bucketName, nil, option.WithoutAuthentication())
-	require.NoError(err)
-
-	testDEK1 := []byte("test DEK")
-	testDEK2 := []byte("more test DEK")
-
-	// request unset value
-	_, err = storage.Get(ctx, "test:input")
 				assert.Error(err)
 
-	// test Put method
+				if tc.unsetError {
-	assert.NoError(storage.Put(ctx, "volume01", testDEK1))
-	assert.NoError(storage.Put(ctx, "volume02", testDEK2))
-
-	// make sure values have been set
-	val, err := storage.Get(ctx, "volume01")
-	assert.NoError(err)
-	assert.Equal(testDEK1, val)
-	val, err = storage.Get(ctx, "volume02")
-	assert.NoError(err)
-	assert.Equal(testDEK2, val)
-
-	_, err = storage.Get(ctx, "invalid:key")
-	assert.Error(err)
 					assert.ErrorIs(err, ErrDEKUnset)
-}
+				} else {
-
+					assert.False(errors.Is(err, ErrDEKUnset))
-func setupEmulator(ctx context.Context, cli *client.Client, imageName string) (container.ContainerCreateCreatedBody, error) {
-	reader, err := cli.ImagePull(ctx, imageName, types.ImagePullOptions{})
-	if err != nil {
-		return container.ContainerCreateCreatedBody{}, err
-	}
-	if _, err := io.Copy(os.Stdout, reader); err != nil {
-		return container.ContainerCreateCreatedBody{}, err
-	}
-	if err := reader.Close(); err != nil {
-		return container.ContainerCreateCreatedBody{}, err
 				}
 
-	// the 3 true statements are necessary to attach later to the container log
+			} else {
-	containerConfig := &container.Config{
+				assert.NoError(err)
-		Image:        storageEmulator,
+				assert.Equal(tc.client.newReaderOutput, out)
-		AttachStdout: true,
+			}
-		AttachStderr: true,
+		})
-		Tty:          true,
+	}
+}
+
+func TestGCPPut(t *testing.T) {
+	someErr := errors.New("error")
+	testCases := map[string]struct {
+		client      *stubGCPStorageAPI
+		unsetError  bool
+		errExpected bool
+	}{
+		"success": {
+			client: &stubGCPStorageAPI{
+				writer: &stubWriteCloser{
+					result: new([]byte),
+				},
+			},
+		},
+		"creating client fails": {
+			client:      &stubGCPStorageAPI{newClientErr: someErr},
+			errExpected: true,
+		},
+		"NewWriter fails": {
+			client: &stubGCPStorageAPI{
+				writer: &stubWriteCloser{
+					result:   new([]byte),
+					writeErr: someErr,
+				},
+			},
+			errExpected: true,
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+
+			client := &GoogleCloudStorage{
+				newClient:  tc.client.stubClientFactory,
+				projectID:  "test",
+				bucketName: "test",
+			}
+			testData := []byte{0x1, 0x2, 0x3}
+
+			err := client.Put(context.Background(), "test-key", testData)
+			if tc.errExpected {
+				assert.Error(err)
+			} else {
+				assert.NoError(err)
+				assert.Equal(testData, *tc.client.writer.result)
+			}
+		})
+	}
+}
+
+func TestGCPCreateContainerOrContinue(t *testing.T) {
+	someErr := errors.New("error")
+	testCases := map[string]struct {
+		client          *stubGCPStorageAPI
+		createNewBucket bool
+		errExpected     bool
+	}{
+		"success": {
+			client: &stubGCPStorageAPI{},
+		},
+		"container does not exist": {
+			client:          &stubGCPStorageAPI{attrsErr: storage.ErrBucketNotExist},
+			createNewBucket: true,
+		},
+		"creating client fails": {
+			client:      &stubGCPStorageAPI{newClientErr: someErr},
+			errExpected: true,
+		},
+		"Attrs fails": {
+			client:      &stubGCPStorageAPI{attrsErr: someErr},
+			errExpected: true,
+		},
+		"CreateBucket fails": {
+			client: &stubGCPStorageAPI{
+				attrsErr:        storage.ErrBucketNotExist,
+				createBucketErr: someErr,
+			},
+			errExpected: true,
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+
+			client := &GoogleCloudStorage{
+				newClient:  tc.client.stubClientFactory,
+				projectID:  "test",
+				bucketName: "test",
+			}
+
+			err := client.createContainerOrContinue(context.Background(), nil)
+			if tc.errExpected {
+				assert.Error(err)
+			} else {
+				assert.NoError(err)
+				if tc.createNewBucket {
+					assert.True(tc.client.createBucketCalled)
+				}
 			}
-	emulator, err := cli.ContainerCreate(ctx, containerConfig, &container.HostConfig{NetworkMode: container.NetworkMode("host"), AutoRemove: true}, nil, nil, "google-cloud-storage-test")
-	if err != nil {
-		return emulator, err
-	}
-	err = cli.ContainerStart(ctx, emulator.ID, types.ContainerStartOptions{})
-	if err != nil {
-		return emulator, err
-	}
-
-	logs, err := cli.ContainerLogs(ctx, emulator.ID, types.ContainerLogsOptions{
-		ShowStdout: true,
-		Follow:     true,
 		})
-	if err != nil {
-		return emulator, err
 	}
-	go io.Copy(os.Stdout, logs)
-	return emulator, nil
 }