Add test for wrapped AWS errors

internal/filestore/file_handler_test.go

@@ -293,29 +293,62 @@ func TestSaveFile(t *testing.T) {
 }
 
 func TestSaveFileWithS3WorkhorseClient(t *testing.T) {
-	s3Creds, s3Config, sess, ts := test.SetupS3(t, "")
-	defer ts.Close()
-
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	remoteObject := "tmp/test-file/1"
-	opts := filestore.SaveFileOpts{
-		RemoteID:           "test-file",
-		Deadline:           testDeadline(),
-		UseWorkhorseClient: true,
-		RemoteTempObjectID: remoteObject,
-		ObjectStorageConfig: filestore.ObjectStorageConfig{
-			Provider:      "AWS",
-			S3Credentials: s3Creds,
-			S3Config:      s3Config,
+	tests := []struct {
+		name        string
+		objectSize  int64
+		maxSize     int64
+		expectedErr error
+	}{
+		{
+			name:       "known size with no limit",
+			objectSize: test.ObjectSize,
+		},
+		{
+			name:       "unknown size with no limit",
+			objectSize: -1,
+		},
+		{
+			name:        "unknown object size with limit",
+			objectSize:  -1,
+			maxSize:     test.ObjectSize - 1,
+			expectedErr: filestore.ErrEntityTooLarge,
 		},
 	}
 
-	_, err := filestore.SaveFileFromReader(ctx, strings.NewReader(test.ObjectContent), test.ObjectSize, &opts)
-	require.NoError(t, err)
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+
+			s3Creds, s3Config, sess, ts := test.SetupS3(t, "")
+			defer ts.Close()
 
-	test.S3ObjectExists(t, sess, s3Config, remoteObject, test.ObjectContent)
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			remoteObject := "tmp/test-file/1"
+			opts := filestore.SaveFileOpts{
+				RemoteID:           "test-file",
+				Deadline:           testDeadline(),
+				UseWorkhorseClient: true,
+				RemoteTempObjectID: remoteObject,
+				ObjectStorageConfig: filestore.ObjectStorageConfig{
+					Provider:      "AWS",
+					S3Credentials: s3Creds,
+					S3Config:      s3Config,
+				},
+				MaximumSize: tc.maxSize,
+			}
+
+			_, err := filestore.SaveFileFromReader(ctx, strings.NewReader(test.ObjectContent), tc.objectSize, &opts)
+
+			if tc.expectedErr == nil {
+				require.NoError(t, err)
+				test.S3ObjectExists(t, sess, s3Config, remoteObject, test.ObjectContent)
+			} else {
+				require.Equal(t, tc.expectedErr, err)
+				test.S3ObjectDoesNotExist(t, sess, s3Config, remoteObject)
+			}
+		})
+	}
 }
 
 func TestSaveFileWithAzureWorkhorseClient(t *testing.T) {

internal/objectstore/s3_object_test.go

@@ -3,6 +3,7 @@ package objectstore_test
 import (
 	"context"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -11,6 +12,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/s3"
 	"github.com/stretchr/testify/require"
@@ -21,6 +23,15 @@ import (
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/testhelper"
 )
 
+type failedReader struct {
+	io.Reader
+}
+
+func (r *failedReader) Read(p []byte) (int, error) {
+	origErr := fmt.Errorf("entity is too large")
+	return 0, awserr.New("Read", "read failed", origErr)
+}
+
 func TestS3ObjectUpload(t *testing.T) {
 	testCases := []struct {
 		encryption string
@@ -143,3 +154,21 @@ func TestS3ObjectUploadCancel(t *testing.T) {
 	require.Error(t, err)
 	require.Equal(t, "context canceled", err.Error())
 }
+
+func TestS3ObjectUploadLimitReached(t *testing.T) {
+	creds, config, _, ts := test.SetupS3(t, "")
+	defer ts.Close()
+
+	deadline := time.Now().Add(testTimeout)
+	tmpDir, err := ioutil.TempDir("", "workhorse-test-")
+	require.NoError(t, err)
+	defer os.Remove(tmpDir)
+
+	objectName := filepath.Join(tmpDir, "s3-test-data")
+	object, err := objectstore.NewS3Object(objectName, creds, config)
+	require.NoError(t, err)
+
+	_, err = object.Consume(context.Background(), &failedReader{}, deadline)
+	require.Error(t, err)
+	require.Equal(t, "entity is too large", err.Error())
+}