s3store: Pass context to AWS SDK

pkg/s3store/s3store.go

@@ -86,6 +86,7 @@ import (
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/s3"
 )
 
@@ -134,16 +135,16 @@ type S3Store struct {
 }
 
 type S3API interface {
-	PutObject(input *s3.PutObjectInput) (*s3.PutObjectOutput, error)
+	PutObjectWithContext(ctx context.Context, input *s3.PutObjectInput, opt ...request.Option) (*s3.PutObjectOutput, error)
-	ListParts(input *s3.ListPartsInput) (*s3.ListPartsOutput, error)
+	ListPartsWithContext(ctx context.Context, input *s3.ListPartsInput, opt ...request.Option) (*s3.ListPartsOutput, error)
-	UploadPart(input *s3.UploadPartInput) (*s3.UploadPartOutput, error)
+	UploadPartWithContext(ctx context.Context, input *s3.UploadPartInput, opt ...request.Option) (*s3.UploadPartOutput, error)
-	GetObject(input *s3.GetObjectInput) (*s3.GetObjectOutput, error)
+	GetObjectWithContext(ctx context.Context, input *s3.GetObjectInput, opt ...request.Option) (*s3.GetObjectOutput, error)
-	CreateMultipartUpload(input *s3.CreateMultipartUploadInput) (*s3.CreateMultipartUploadOutput, error)
+	CreateMultipartUploadWithContext(ctx context.Context, input *s3.CreateMultipartUploadInput, opt ...request.Option) (*s3.CreateMultipartUploadOutput, error)
-	AbortMultipartUpload(input *s3.AbortMultipartUploadInput) (*s3.AbortMultipartUploadOutput, error)
+	AbortMultipartUploadWithContext(ctx context.Context, input *s3.AbortMultipartUploadInput, opt ...request.Option) (*s3.AbortMultipartUploadOutput, error)
-	DeleteObject(input *s3.DeleteObjectInput) (*s3.DeleteObjectOutput, error)
+	DeleteObjectWithContext(ctx context.Context, input *s3.DeleteObjectInput, opt ...request.Option) (*s3.DeleteObjectOutput, error)
-	DeleteObjects(input *s3.DeleteObjectsInput) (*s3.DeleteObjectsOutput, error)
+	DeleteObjectsWithContext(ctx context.Context, input *s3.DeleteObjectsInput, opt ...request.Option) (*s3.DeleteObjectsOutput, error)
-	CompleteMultipartUpload(input *s3.CompleteMultipartUploadInput) (*s3.CompleteMultipartUploadOutput, error)
+	CompleteMultipartUploadWithContext(ctx context.Context, input *s3.CompleteMultipartUploadInput, opt ...request.Option) (*s3.CompleteMultipartUploadOutput, error)
-	UploadPartCopy(input *s3.UploadPartCopyInput) (*s3.UploadPartCopyOutput, error)
+	UploadPartCopyWithContext(ctx context.Context, input *s3.UploadPartCopyInput, opt ...request.Option) (*s3.UploadPartCopyOutput, error)
 }
 
 // New constructs a new storage using the supplied bucket and service object.
@@ -196,7 +197,7 @@ func (store S3Store) NewUpload(ctx context.Context, info handler.FileInfo) (hand
 	}
 
 	// Create the actual multipart upload
-	res, err := store.Service.CreateMultipartUpload(&s3.CreateMultipartUploadInput{
+	res, err := store.Service.CreateMultipartUploadWithContext(ctx, &s3.CreateMultipartUploadInput{
 		Bucket:   aws.String(store.Bucket),
 		Key:      store.keyWithPrefix(uploadId),
 		Metadata: metadata,
@@ -214,7 +215,7 @@ func (store S3Store) NewUpload(ctx context.Context, info handler.FileInfo) (hand
 		"Key":    *store.keyWithPrefix(uploadId),
 	}
 
-	err = store.writeInfo(uploadId, info)
+	err = store.writeInfo(ctx, uploadId, info)
 	if err != nil {
 		return nil, fmt.Errorf("s3store: unable to create info file:\n%s", err)
 	}
@@ -234,14 +235,14 @@ func (store S3Store) AsLengthDeclarableUpload(upload handler.Upload) handler.Len
 	return upload.(*s3Upload)
 }
 
-func (store S3Store) writeInfo(uploadId string, info handler.FileInfo) error {
+func (store S3Store) writeInfo(ctx context.Context, uploadId string, info handler.FileInfo) error {
 	infoJson, err := json.Marshal(info)
 	if err != nil {
 		return err
 	}
 
 	// Create object on S3 containing information about the file
-	_, err = store.Service.PutObject(&s3.PutObjectInput{
+	_, err = store.Service.PutObjectWithContext(ctx, &s3.PutObjectInput{
 		Bucket:        aws.String(store.Bucket),
 		Key:           store.keyWithPrefix(uploadId + ".info"),
 		Body:          bytes.NewReader(infoJson),
@@ -271,7 +272,7 @@ func (upload s3Upload) WriteChunk(ctx context.Context, offset int64, src io.Read
 	}
 
 	// Get number of parts to generate next number
-	parts, err := store.listAllParts(id)
+	parts, err := store.listAllParts(ctx, id)
 	if err != nil {
 		return 0, err
 	}
@@ -279,7 +280,7 @@ func (upload s3Upload) WriteChunk(ctx context.Context, offset int64, src io.Read
 	numParts := len(parts)
 	nextPartNum := int64(numParts + 1)
 
-	incompletePartFile, incompletePartSize, err := store.downloadIncompletePartForUpload(uploadId)
+	incompletePartFile, incompletePartSize, err := store.downloadIncompletePartForUpload(ctx, uploadId)
 	if err != nil {
 		return 0, err
 	}
@@ -287,7 +288,7 @@ func (upload s3Upload) WriteChunk(ctx context.Context, offset int64, src io.Read
 		defer os.Remove(incompletePartFile.Name())
 		defer incompletePartFile.Close()
 
-		if err := store.deleteIncompletePartForUpload(uploadId); err != nil {
+		if err := store.deleteIncompletePartForUpload(ctx, uploadId); err != nil {
 			return 0, err
 		}
 
@@ -329,7 +330,7 @@ func (upload s3Upload) WriteChunk(ctx context.Context, offset int64, src io.Read
 
 		isFinalChunk := !info.SizeIsDeferred && (size == (offset-incompletePartSize)+n)
 		if n >= store.MinPartSize || isFinalChunk {
-			_, err = store.Service.UploadPart(&s3.UploadPartInput{
+			_, err = store.Service.UploadPartWithContext(ctx, &s3.UploadPartInput{
 				Bucket:     aws.String(store.Bucket),
 				Key:        store.keyWithPrefix(uploadId),
 				UploadId:   aws.String(multipartId),
@@ -340,7 +341,7 @@ func (upload s3Upload) WriteChunk(ctx context.Context, offset int64, src io.Read
 				return bytesUploaded, err
 			}
 		} else {
-			if err := store.putIncompletePartForUpload(uploadId, file); err != nil {
+			if err := store.putIncompletePartForUpload(ctx, uploadId, file); err != nil {
 				return bytesUploaded, err
 			}
 
@@ -361,7 +362,7 @@ func (upload s3Upload) GetInfo(ctx context.Context) (info handler.FileInfo, err
 	uploadId, _ := splitIds(id)
 
 	// Get file info stored in separate object
-	res, err := store.Service.GetObject(&s3.GetObjectInput{
+	res, err := store.Service.GetObjectWithContext(ctx, &s3.GetObjectInput{
 		Bucket: aws.String(store.Bucket),
 		Key:    store.keyWithPrefix(uploadId + ".info"),
 	})
@@ -378,7 +379,7 @@ func (upload s3Upload) GetInfo(ctx context.Context) (info handler.FileInfo, err
 	}
 
 	// Get uploaded parts and their offset
-	parts, err := store.listAllParts(id)
+	parts, err := store.listAllParts(ctx, id)
 	if err != nil {
 		// Check if the error is caused by the upload not being found. This happens
 		// when the multipart upload has already been completed or aborted. Since
@@ -398,7 +399,7 @@ func (upload s3Upload) GetInfo(ctx context.Context) (info handler.FileInfo, err
 		offset += *part.Size
 	}
 
-	incompletePartObject, err := store.getIncompletePartForUpload(uploadId)
+	incompletePartObject, err := store.getIncompletePartForUpload(ctx, uploadId)
 	if err != nil {
 		return info, err
 	}
@@ -418,7 +419,7 @@ func (upload s3Upload) GetReader(ctx context.Context) (io.Reader, error) {
 	uploadId, multipartId := splitIds(id)
 
 	// Attempt to get upload content
-	res, err := store.Service.GetObject(&s3.GetObjectInput{
+	res, err := store.Service.GetObjectWithContext(ctx, &s3.GetObjectInput{
 		Bucket: aws.String(store.Bucket),
 		Key:    store.keyWithPrefix(uploadId),
 	})
@@ -436,7 +437,7 @@ func (upload s3Upload) GetReader(ctx context.Context) (io.Reader, error) {
 
 	// Test whether the multipart upload exists to find out if the upload
 	// never existsted or just has not been finished yet
-	_, err = store.Service.ListParts(&s3.ListPartsInput{
+	_, err = store.Service.ListPartsWithContext(ctx, &s3.ListPartsInput{
 		Bucket:   aws.String(store.Bucket),
 		Key:      store.keyWithPrefix(uploadId),
 		UploadId: aws.String(multipartId),
@@ -467,7 +468,7 @@ func (upload s3Upload) Terminate(ctx context.Context) error {
 		defer wg.Done()
 
 		// Abort the multipart upload
-		_, err := store.Service.AbortMultipartUpload(&s3.AbortMultipartUploadInput{
+		_, err := store.Service.AbortMultipartUploadWithContext(ctx, &s3.AbortMultipartUploadInput{
 			Bucket:   aws.String(store.Bucket),
 			Key:      store.keyWithPrefix(uploadId),
 			UploadId: aws.String(multipartId),
@@ -481,7 +482,7 @@ func (upload s3Upload) Terminate(ctx context.Context) error {
 		defer wg.Done()
 
 		// Delete the info and content files
-		res, err := store.Service.DeleteObjects(&s3.DeleteObjectsInput{
+		res, err := store.Service.DeleteObjectsWithContext(ctx, &s3.DeleteObjectsInput{
 			Bucket: aws.String(store.Bucket),
 			Delete: &s3.Delete{
 				Objects: []*s3.ObjectIdentifier{
@@ -526,7 +527,7 @@ func (upload s3Upload) FinishUpload(ctx context.Context) error {
 	uploadId, multipartId := splitIds(id)
 
 	// Get uploaded parts
-	parts, err := store.listAllParts(id)
+	parts, err := store.listAllParts(ctx, id)
 	if err != nil {
 		return err
 	}
@@ -542,7 +543,7 @@ func (upload s3Upload) FinishUpload(ctx context.Context) error {
 		}
 	}
 
-	_, err = store.Service.CompleteMultipartUpload(&s3.CompleteMultipartUploadInput{
+	_, err = store.Service.CompleteMultipartUploadWithContext(ctx, &s3.CompleteMultipartUploadInput{
 		Bucket:   aws.String(store.Bucket),
 		Key:      store.keyWithPrefix(uploadId),
 		UploadId: aws.String(multipartId),
@@ -569,7 +570,7 @@ func (store S3Store) ConcatUploads(ctx context.Context, dest string, partialUplo
 
 			partialUploadId, _ := splitIds(partialId)
 
-			_, err := store.Service.UploadPartCopy(&s3.UploadPartCopyInput{
+			_, err := store.Service.UploadPartCopyWithContext(ctx, &s3.UploadPartCopyInput{
 				Bucket:   aws.String(store.Bucket),
 				Key:      store.keyWithPrefix(uploadId),
 				UploadId: aws.String(multipartId),
@@ -609,16 +610,16 @@ func (upload s3Upload) DeclareLength(ctx context.Context, length int64) error {
 	info.Size = length
 	info.SizeIsDeferred = false
 
-	return store.writeInfo(uploadId, info)
+	return store.writeInfo(ctx, uploadId, info)
 }
 
-func (store S3Store) listAllParts(id string) (parts []*s3.Part, err error) {
+func (store S3Store) listAllParts(ctx context.Context, id string) (parts []*s3.Part, err error) {
 	uploadId, multipartId := splitIds(id)
 
 	partMarker := int64(0)
 	for {
 		// Get uploaded parts
-		listPtr, err := store.Service.ListParts(&s3.ListPartsInput{
+		listPtr, err := store.Service.ListPartsWithContext(ctx, &s3.ListPartsInput{
 			Bucket:           aws.String(store.Bucket),
 			Key:              store.keyWithPrefix(uploadId),
 			UploadId:         aws.String(multipartId),
@@ -639,8 +640,8 @@ func (store S3Store) listAllParts(id string) (parts []*s3.Part, err error) {
 	return parts, nil
 }
 
-func (store S3Store) downloadIncompletePartForUpload(uploadId string) (*os.File, int64, error) {
+func (store S3Store) downloadIncompletePartForUpload(ctx context.Context, uploadId string) (*os.File, int64, error) {
-	incompleteUploadObject, err := store.getIncompletePartForUpload(uploadId)
+	incompleteUploadObject, err := store.getIncompletePartForUpload(ctx, uploadId)
 	if err != nil {
 		return nil, 0, err
 	}
@@ -671,8 +672,8 @@ func (store S3Store) downloadIncompletePartForUpload(uploadId string) (*os.File,
 	return partFile, n, nil
 }
 
-func (store S3Store) getIncompletePartForUpload(uploadId string) (*s3.GetObjectOutput, error) {
+func (store S3Store) getIncompletePartForUpload(ctx context.Context, uploadId string) (*s3.GetObjectOutput, error) {
-	obj, err := store.Service.GetObject(&s3.GetObjectInput{
+	obj, err := store.Service.GetObjectWithContext(ctx, &s3.GetObjectInput{
 		Bucket: aws.String(store.Bucket),
 		Key:    store.keyWithPrefix(uploadId + ".part"),
 	})
@@ -684,8 +685,8 @@ func (store S3Store) getIncompletePartForUpload(uploadId string) (*s3.GetObjectO
 	return obj, err
 }
 
-func (store S3Store) putIncompletePartForUpload(uploadId string, r io.ReadSeeker) error {
+func (store S3Store) putIncompletePartForUpload(ctx context.Context, uploadId string, r io.ReadSeeker) error {
-	_, err := store.Service.PutObject(&s3.PutObjectInput{
+	_, err := store.Service.PutObjectWithContext(ctx, &s3.PutObjectInput{
 		Bucket: aws.String(store.Bucket),
 		Key:    store.keyWithPrefix(uploadId + ".part"),
 		Body:   r,
@@ -693,8 +694,8 @@ func (store S3Store) putIncompletePartForUpload(uploadId string, r io.ReadSeeker
 	return err
 }
 
-func (store S3Store) deleteIncompletePartForUpload(uploadId string) error {
+func (store S3Store) deleteIncompletePartForUpload(ctx context.Context, uploadId string) error {
-	_, err := store.Service.DeleteObject(&s3.DeleteObjectInput{
+	_, err := store.Service.DeleteObjectWithContext(ctx, &s3.DeleteObjectInput{
 		Bucket: aws.String(store.Bucket),
 		Key:    store.keyWithPrefix(uploadId + ".part"),
 	})

pkg/s3store/s3store_mock_test.go

@@ -1,140 +1,236 @@
-// Automatically generated by MockGen. DO NOT EDIT!
+// Code generated by MockGen. DO NOT EDIT.
 // Source: github.com/tus/tusd/pkg/s3store (interfaces: S3API)
 
+// Package s3store is a generated GoMock package.
 package s3store
 
 import (
+	context "context"
+	request "github.com/aws/aws-sdk-go/aws/request"
 	s3 "github.com/aws/aws-sdk-go/service/s3"
 	gomock "github.com/golang/mock/gomock"
+	reflect "reflect"
 )
 
-// Mock of S3API interface
+// MockS3API is a mock of S3API interface
 type MockS3API struct {
 	ctrl     *gomock.Controller
-	recorder *_MockS3APIRecorder
+	recorder *MockS3APIMockRecorder
 }
 
-// Recorder for MockS3API (not exported)
+// MockS3APIMockRecorder is the mock recorder for MockS3API
-type _MockS3APIRecorder struct {
+type MockS3APIMockRecorder struct {
 	mock *MockS3API
 }
 
+// NewMockS3API creates a new mock instance
 func NewMockS3API(ctrl *gomock.Controller) *MockS3API {
 	mock := &MockS3API{ctrl: ctrl}
-	mock.recorder = &_MockS3APIRecorder{mock}
+	mock.recorder = &MockS3APIMockRecorder{mock}
 	return mock
 }
 
-func (_m *MockS3API) EXPECT() *_MockS3APIRecorder {
+// EXPECT returns an object that allows the caller to indicate expected use
-	return _m.recorder
+func (m *MockS3API) EXPECT() *MockS3APIMockRecorder {
+	return m.recorder
 }
 
-func (_m *MockS3API) AbortMultipartUpload(_param0 *s3.AbortMultipartUploadInput) (*s3.AbortMultipartUploadOutput, error) {
+// AbortMultipartUploadWithContext mocks base method
-	ret := _m.ctrl.Call(_m, "AbortMultipartUpload", _param0)
+func (m *MockS3API) AbortMultipartUploadWithContext(arg0 context.Context, arg1 *s3.AbortMultipartUploadInput, arg2 ...request.Option) (*s3.AbortMultipartUploadOutput, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "AbortMultipartUploadWithContext", varargs...)
 	ret0, _ := ret[0].(*s3.AbortMultipartUploadOutput)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockS3APIRecorder) AbortMultipartUpload(arg0 interface{}) *gomock.Call {
+// AbortMultipartUploadWithContext indicates an expected call of AbortMultipartUploadWithContext
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "AbortMultipartUpload", arg0)
+func (mr *MockS3APIMockRecorder) AbortMultipartUploadWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AbortMultipartUploadWithContext", reflect.TypeOf((*MockS3API)(nil).AbortMultipartUploadWithContext), varargs...)
 }
 
-func (_m *MockS3API) CompleteMultipartUpload(_param0 *s3.CompleteMultipartUploadInput) (*s3.CompleteMultipartUploadOutput, error) {
+// CompleteMultipartUploadWithContext mocks base method
-	ret := _m.ctrl.Call(_m, "CompleteMultipartUpload", _param0)
+func (m *MockS3API) CompleteMultipartUploadWithContext(arg0 context.Context, arg1 *s3.CompleteMultipartUploadInput, arg2 ...request.Option) (*s3.CompleteMultipartUploadOutput, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "CompleteMultipartUploadWithContext", varargs...)
 	ret0, _ := ret[0].(*s3.CompleteMultipartUploadOutput)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockS3APIRecorder) CompleteMultipartUpload(arg0 interface{}) *gomock.Call {
+// CompleteMultipartUploadWithContext indicates an expected call of CompleteMultipartUploadWithContext
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "CompleteMultipartUpload", arg0)
+func (mr *MockS3APIMockRecorder) CompleteMultipartUploadWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CompleteMultipartUploadWithContext", reflect.TypeOf((*MockS3API)(nil).CompleteMultipartUploadWithContext), varargs...)
 }
 
-func (_m *MockS3API) CreateMultipartUpload(_param0 *s3.CreateMultipartUploadInput) (*s3.CreateMultipartUploadOutput, error) {
+// CreateMultipartUploadWithContext mocks base method
-	ret := _m.ctrl.Call(_m, "CreateMultipartUpload", _param0)
+func (m *MockS3API) CreateMultipartUploadWithContext(arg0 context.Context, arg1 *s3.CreateMultipartUploadInput, arg2 ...request.Option) (*s3.CreateMultipartUploadOutput, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "CreateMultipartUploadWithContext", varargs...)
 	ret0, _ := ret[0].(*s3.CreateMultipartUploadOutput)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockS3APIRecorder) CreateMultipartUpload(arg0 interface{}) *gomock.Call {
+// CreateMultipartUploadWithContext indicates an expected call of CreateMultipartUploadWithContext
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "CreateMultipartUpload", arg0)
+func (mr *MockS3APIMockRecorder) CreateMultipartUploadWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateMultipartUploadWithContext", reflect.TypeOf((*MockS3API)(nil).CreateMultipartUploadWithContext), varargs...)
 }
 
-func (_m *MockS3API) DeleteObject(_param0 *s3.DeleteObjectInput) (*s3.DeleteObjectOutput, error) {
+// DeleteObjectWithContext mocks base method
-	ret := _m.ctrl.Call(_m, "DeleteObject", _param0)
+func (m *MockS3API) DeleteObjectWithContext(arg0 context.Context, arg1 *s3.DeleteObjectInput, arg2 ...request.Option) (*s3.DeleteObjectOutput, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "DeleteObjectWithContext", varargs...)
 	ret0, _ := ret[0].(*s3.DeleteObjectOutput)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockS3APIRecorder) DeleteObject(arg0 interface{}) *gomock.Call {
+// DeleteObjectWithContext indicates an expected call of DeleteObjectWithContext
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "DeleteObject", arg0)
+func (mr *MockS3APIMockRecorder) DeleteObjectWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteObjectWithContext", reflect.TypeOf((*MockS3API)(nil).DeleteObjectWithContext), varargs...)
 }
 
-func (_m *MockS3API) DeleteObjects(_param0 *s3.DeleteObjectsInput) (*s3.DeleteObjectsOutput, error) {
+// DeleteObjectsWithContext mocks base method
-	ret := _m.ctrl.Call(_m, "DeleteObjects", _param0)
+func (m *MockS3API) DeleteObjectsWithContext(arg0 context.Context, arg1 *s3.DeleteObjectsInput, arg2 ...request.Option) (*s3.DeleteObjectsOutput, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "DeleteObjectsWithContext", varargs...)
 	ret0, _ := ret[0].(*s3.DeleteObjectsOutput)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockS3APIRecorder) DeleteObjects(arg0 interface{}) *gomock.Call {
+// DeleteObjectsWithContext indicates an expected call of DeleteObjectsWithContext
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "DeleteObjects", arg0)
+func (mr *MockS3APIMockRecorder) DeleteObjectsWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteObjectsWithContext", reflect.TypeOf((*MockS3API)(nil).DeleteObjectsWithContext), varargs...)
 }
 
-func (_m *MockS3API) GetObject(_param0 *s3.GetObjectInput) (*s3.GetObjectOutput, error) {
+// GetObjectWithContext mocks base method
-	ret := _m.ctrl.Call(_m, "GetObject", _param0)
+func (m *MockS3API) GetObjectWithContext(arg0 context.Context, arg1 *s3.GetObjectInput, arg2 ...request.Option) (*s3.GetObjectOutput, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "GetObjectWithContext", varargs...)
 	ret0, _ := ret[0].(*s3.GetObjectOutput)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockS3APIRecorder) GetObject(arg0 interface{}) *gomock.Call {
+// GetObjectWithContext indicates an expected call of GetObjectWithContext
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "GetObject", arg0)
+func (mr *MockS3APIMockRecorder) GetObjectWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetObjectWithContext", reflect.TypeOf((*MockS3API)(nil).GetObjectWithContext), varargs...)
 }
 
-func (_m *MockS3API) ListParts(_param0 *s3.ListPartsInput) (*s3.ListPartsOutput, error) {
+// ListPartsWithContext mocks base method
-	ret := _m.ctrl.Call(_m, "ListParts", _param0)
+func (m *MockS3API) ListPartsWithContext(arg0 context.Context, arg1 *s3.ListPartsInput, arg2 ...request.Option) (*s3.ListPartsOutput, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "ListPartsWithContext", varargs...)
 	ret0, _ := ret[0].(*s3.ListPartsOutput)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockS3APIRecorder) ListParts(arg0 interface{}) *gomock.Call {
+// ListPartsWithContext indicates an expected call of ListPartsWithContext
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ListParts", arg0)
+func (mr *MockS3APIMockRecorder) ListPartsWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListPartsWithContext", reflect.TypeOf((*MockS3API)(nil).ListPartsWithContext), varargs...)
 }
 
-func (_m *MockS3API) PutObject(_param0 *s3.PutObjectInput) (*s3.PutObjectOutput, error) {
+// PutObjectWithContext mocks base method
-	ret := _m.ctrl.Call(_m, "PutObject", _param0)
+func (m *MockS3API) PutObjectWithContext(arg0 context.Context, arg1 *s3.PutObjectInput, arg2 ...request.Option) (*s3.PutObjectOutput, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "PutObjectWithContext", varargs...)
 	ret0, _ := ret[0].(*s3.PutObjectOutput)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockS3APIRecorder) PutObject(arg0 interface{}) *gomock.Call {
+// PutObjectWithContext indicates an expected call of PutObjectWithContext
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "PutObject", arg0)
+func (mr *MockS3APIMockRecorder) PutObjectWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PutObjectWithContext", reflect.TypeOf((*MockS3API)(nil).PutObjectWithContext), varargs...)
 }
 
-func (_m *MockS3API) UploadPart(_param0 *s3.UploadPartInput) (*s3.UploadPartOutput, error) {
+// UploadPartCopyWithContext mocks base method
-	ret := _m.ctrl.Call(_m, "UploadPart", _param0)
+func (m *MockS3API) UploadPartCopyWithContext(arg0 context.Context, arg1 *s3.UploadPartCopyInput, arg2 ...request.Option) (*s3.UploadPartCopyOutput, error) {
-	ret0, _ := ret[0].(*s3.UploadPartOutput)
+	m.ctrl.T.Helper()
-	ret1, _ := ret[1].(error)
+	varargs := []interface{}{arg0, arg1}
-	return ret0, ret1
+	for _, a := range arg2 {
-}
+		varargs = append(varargs, a)
-
+	}
-func (_mr *_MockS3APIRecorder) UploadPart(arg0 interface{}) *gomock.Call {
+	ret := m.ctrl.Call(m, "UploadPartCopyWithContext", varargs...)
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "UploadPart", arg0)
-}
-
-func (_m *MockS3API) UploadPartCopy(_param0 *s3.UploadPartCopyInput) (*s3.UploadPartCopyOutput, error) {
-	ret := _m.ctrl.Call(_m, "UploadPartCopy", _param0)
 	ret0, _ := ret[0].(*s3.UploadPartCopyOutput)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockS3APIRecorder) UploadPartCopy(arg0 interface{}) *gomock.Call {
+// UploadPartCopyWithContext indicates an expected call of UploadPartCopyWithContext
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "UploadPartCopy", arg0)
+func (mr *MockS3APIMockRecorder) UploadPartCopyWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UploadPartCopyWithContext", reflect.TypeOf((*MockS3API)(nil).UploadPartCopyWithContext), varargs...)
+}
+
+// UploadPartWithContext mocks base method
+func (m *MockS3API) UploadPartWithContext(arg0 context.Context, arg1 *s3.UploadPartInput, arg2 ...request.Option) (*s3.UploadPartOutput, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "UploadPartWithContext", varargs...)
+	ret0, _ := ret[0].(*s3.UploadPartOutput)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// UploadPartWithContext indicates an expected call of UploadPartWithContext
+func (mr *MockS3APIMockRecorder) UploadPartWithContext(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UploadPartWithContext", reflect.TypeOf((*MockS3API)(nil).UploadPartWithContext), varargs...)
 }

pkg/s3store/s3store_test.go

@@ -40,7 +40,7 @@ func TestNewUpload(t *testing.T) {
 	s2 := "men???hi"
 
 	gomock.InOrder(
-		s3obj.EXPECT().CreateMultipartUpload(&s3.CreateMultipartUploadInput{
+		s3obj.EXPECT().CreateMultipartUploadWithContext(context.Background(), &s3.CreateMultipartUploadInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId"),
 			Metadata: map[string]*string{
@@ -50,7 +50,7 @@ func TestNewUpload(t *testing.T) {
 		}).Return(&s3.CreateMultipartUploadOutput{
 			UploadId: aws.String("multipartId"),
 		}, nil),
-		s3obj.EXPECT().PutObject(&s3.PutObjectInput{
+		s3obj.EXPECT().PutObjectWithContext(context.Background(), &s3.PutObjectInput{
 			Bucket:        aws.String("bucket"),
 			Key:           aws.String("uploadId.info"),
 			Body:          bytes.NewReader([]byte(`{"ID":"uploadId+multipartId","Size":500,"SizeIsDeferred":false,"Offset":0,"MetaData":{"bar":"menü\r\nhi","foo":"hello"},"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":{"Bucket":"bucket","Key":"uploadId","Type":"s3store"}}`)),
@@ -88,7 +88,7 @@ func TestNewUploadWithObjectPrefix(t *testing.T) {
 	s2 := "men?"
 
 	gomock.InOrder(
-		s3obj.EXPECT().CreateMultipartUpload(&s3.CreateMultipartUploadInput{
+		s3obj.EXPECT().CreateMultipartUploadWithContext(context.Background(), &s3.CreateMultipartUploadInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("my/uploaded/files/uploadId"),
 			Metadata: map[string]*string{
@@ -98,7 +98,7 @@ func TestNewUploadWithObjectPrefix(t *testing.T) {
 		}).Return(&s3.CreateMultipartUploadOutput{
 			UploadId: aws.String("multipartId"),
 		}, nil),
-		s3obj.EXPECT().PutObject(&s3.PutObjectInput{
+		s3obj.EXPECT().PutObjectWithContext(context.Background(), &s3.PutObjectInput{
 			Bucket:        aws.String("bucket"),
 			Key:           aws.String("my/uploaded/files/uploadId.info"),
 			Body:          bytes.NewReader([]byte(`{"ID":"uploadId+multipartId","Size":500,"SizeIsDeferred":false,"Offset":0,"MetaData":{"bar":"menü","foo":"hello"},"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":{"Bucket":"bucket","Key":"my/uploaded/files/uploadId","Type":"s3store"}}`)),
@@ -150,7 +150,7 @@ func TestGetInfoNotFound(t *testing.T) {
 	s3obj := NewMockS3API(mockCtrl)
 	store := New("bucket", s3obj)
 
-	s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+	s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 		Bucket: aws.String("bucket"),
 		Key:    aws.String("uploadId.info"),
 	}).Return(nil, awserr.New("NoSuchKey", "The specified key does not exist.", nil))
@@ -171,13 +171,13 @@ func TestGetInfo(t *testing.T) {
 	store := New("bucket", s3obj)
 
 	gomock.InOrder(
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.info"),
 		}).Return(&s3.GetObjectOutput{
 			Body: ioutil.NopCloser(bytes.NewReader([]byte(`{"ID":"uploadId+multipartId","Size":500,"Offset":0,"MetaData":{"bar":"menü","foo":"hello"},"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":{"Bucket":"bucket","Key":"my/uploaded/files/uploadId","Type":"s3store"}}`))),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -194,7 +194,7 @@ func TestGetInfo(t *testing.T) {
 			NextPartNumberMarker: aws.Int64(2),
 			IsTruncated:          aws.Bool(true),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -206,7 +206,7 @@ func TestGetInfo(t *testing.T) {
 				},
 			},
 		}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{}, awserr.New("NoSuchKey", "Not found", nil)),
@@ -236,19 +236,19 @@ func TestGetInfoWithIncompletePart(t *testing.T) {
 	store := New("bucket", s3obj)
 
 	gomock.InOrder(
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.info"),
 		}).Return(&s3.GetObjectOutput{
 			Body: ioutil.NopCloser(bytes.NewReader([]byte(`{"ID":"uploadId+multipartId","Size":500,"Offset":0,"MetaData":{},"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":null}`))),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
 			PartNumberMarker: aws.Int64(0),
 		}).Return(&s3.ListPartsOutput{Parts: []*s3.Part{}}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{
@@ -275,13 +275,13 @@ func TestGetInfoFinished(t *testing.T) {
 	store := New("bucket", s3obj)
 
 	gomock.InOrder(
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.info"),
 		}).Return(&s3.GetObjectOutput{
 			Body: ioutil.NopCloser(bytes.NewReader([]byte(`{"ID":"uploadId","Size":500,"Offset":0,"MetaData":null,"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":null}`))),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -306,7 +306,7 @@ func TestGetReader(t *testing.T) {
 	s3obj := NewMockS3API(mockCtrl)
 	store := New("bucket", s3obj)
 
-	s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+	s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 		Bucket: aws.String("bucket"),
 		Key:    aws.String("uploadId"),
 	}).Return(&s3.GetObjectOutput{
@@ -330,11 +330,11 @@ func TestGetReaderNotFound(t *testing.T) {
 	store := New("bucket", s3obj)
 
 	gomock.InOrder(
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId"),
 		}).Return(nil, awserr.New("NoSuchKey", "The specified key does not exist.", nil)),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:   aws.String("bucket"),
 			Key:      aws.String("uploadId"),
 			UploadId: aws.String("multipartId"),
@@ -359,11 +359,11 @@ func TestGetReaderNotFinished(t *testing.T) {
 	store := New("bucket", s3obj)
 
 	gomock.InOrder(
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId"),
 		}).Return(nil, awserr.New("NoSuchKey", "The specified key does not exist.", nil)),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:   aws.String("bucket"),
 			Key:      aws.String("uploadId"),
 			UploadId: aws.String("multipartId"),
@@ -390,13 +390,13 @@ func TestDeclareLength(t *testing.T) {
 	store := New("bucket", s3obj)
 
 	gomock.InOrder(
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.info"),
 		}).Return(&s3.GetObjectOutput{
 			Body: ioutil.NopCloser(bytes.NewReader([]byte(`{"ID":"uploadId+multipartId","Size":0,"SizeIsDeferred":true,"Offset":0,"MetaData":{},"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":{"Bucket":"bucket","Key":"uploadId","Type":"s3store"}}`))),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -404,11 +404,11 @@ func TestDeclareLength(t *testing.T) {
 		}).Return(&s3.ListPartsOutput{
 			Parts: []*s3.Part{},
 		}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(nil, awserr.New("NotFound", "Not Found", nil)),
-		s3obj.EXPECT().PutObject(&s3.PutObjectInput{
+		s3obj.EXPECT().PutObjectWithContext(context.Background(), &s3.PutObjectInput{
 			Bucket:        aws.String("bucket"),
 			Key:           aws.String("uploadId.info"),
 			Body:          bytes.NewReader([]byte(`{"ID":"uploadId+multipartId","Size":500,"SizeIsDeferred":false,"Offset":0,"MetaData":{},"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":{"Bucket":"bucket","Key":"uploadId","Type":"s3store"}}`)),
@@ -432,7 +432,7 @@ func TestFinishUpload(t *testing.T) {
 	store := New("bucket", s3obj)
 
 	gomock.InOrder(
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -453,7 +453,7 @@ func TestFinishUpload(t *testing.T) {
 			NextPartNumberMarker: aws.Int64(2),
 			IsTruncated:          aws.Bool(true),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -467,7 +467,7 @@ func TestFinishUpload(t *testing.T) {
 				},
 			},
 		}, nil),
-		s3obj.EXPECT().CompleteMultipartUpload(&s3.CompleteMultipartUploadInput{
+		s3obj.EXPECT().CompleteMultipartUploadWithContext(context.Background(), &s3.CompleteMultipartUploadInput{
 			Bucket:   aws.String("bucket"),
 			Key:      aws.String("uploadId"),
 			UploadId: aws.String("multipartId"),
@@ -510,13 +510,13 @@ func TestWriteChunk(t *testing.T) {
 	store.MaxObjectSize = 5 * 1024 * 1024 * 1024 * 1024
 
 	gomock.InOrder(
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.info"),
 		}).Return(&s3.GetObjectOutput{
 			Body: ioutil.NopCloser(bytes.NewReader([]byte(`{"ID":"uploadId","Size":500,"Offset":0,"MetaData":null,"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":null}`))),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -531,11 +531,11 @@ func TestWriteChunk(t *testing.T) {
 				},
 			},
 		}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{}, awserr.New("NoSuchKey", "Not found", nil)),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -550,32 +550,32 @@ func TestWriteChunk(t *testing.T) {
 				},
 			},
 		}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{}, awserr.New("NoSuchKey", "The specified key does not exist.", nil)),
-		s3obj.EXPECT().UploadPart(NewUploadPartInputMatcher(&s3.UploadPartInput{
+		s3obj.EXPECT().UploadPartWithContext(context.Background(), NewUploadPartInputMatcher(&s3.UploadPartInput{
 			Bucket:     aws.String("bucket"),
 			Key:        aws.String("uploadId"),
 			UploadId:   aws.String("multipartId"),
 			PartNumber: aws.Int64(3),
 			Body:       bytes.NewReader([]byte("1234")),
 		})).Return(nil, nil),
-		s3obj.EXPECT().UploadPart(NewUploadPartInputMatcher(&s3.UploadPartInput{
+		s3obj.EXPECT().UploadPartWithContext(context.Background(), NewUploadPartInputMatcher(&s3.UploadPartInput{
 			Bucket:     aws.String("bucket"),
 			Key:        aws.String("uploadId"),
 			UploadId:   aws.String("multipartId"),
 			PartNumber: aws.Int64(4),
 			Body:       bytes.NewReader([]byte("5678")),
 		})).Return(nil, nil),
-		s3obj.EXPECT().UploadPart(NewUploadPartInputMatcher(&s3.UploadPartInput{
+		s3obj.EXPECT().UploadPartWithContext(context.Background(), NewUploadPartInputMatcher(&s3.UploadPartInput{
 			Bucket:     aws.String("bucket"),
 			Key:        aws.String("uploadId"),
 			UploadId:   aws.String("multipartId"),
 			PartNumber: aws.Int64(5),
 			Body:       bytes.NewReader([]byte("90AB")),
 		})).Return(nil, nil),
-		s3obj.EXPECT().PutObject(NewPutObjectInputMatcher(&s3.PutObjectInput{
+		s3obj.EXPECT().PutObjectWithContext(context.Background(), NewPutObjectInputMatcher(&s3.PutObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 			Body:   bytes.NewReader([]byte("CD")),
@@ -606,13 +606,13 @@ func TestWriteChunkWithUnexpectedEOF(t *testing.T) {
 	store.MaxObjectSize = 5 * 1024 * 1024 * 1024 * 1024
 
 	gomock.InOrder(
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.info"),
 		}).Return(&s3.GetObjectOutput{
 			Body: ioutil.NopCloser(bytes.NewReader([]byte(`{"ID":"uploadId","Size":500,"Offset":0,"MetaData":null,"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":null}`))),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -627,11 +627,11 @@ func TestWriteChunkWithUnexpectedEOF(t *testing.T) {
 				},
 			},
 		}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{}, awserr.New("NoSuchKey", "Not found", nil)),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -646,11 +646,11 @@ func TestWriteChunkWithUnexpectedEOF(t *testing.T) {
 				},
 			},
 		}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{}, awserr.New("NoSuchKey", "The specified key does not exist.", nil)),
-		s3obj.EXPECT().PutObject(NewPutObjectInputMatcher(&s3.PutObjectInput{
+		s3obj.EXPECT().PutObjectWithContext(context.Background(), NewPutObjectInputMatcher(&s3.PutObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 			Body:   bytes.NewReader([]byte("1234567890ABCD")),
@@ -681,13 +681,13 @@ func TestWriteChunkWriteIncompletePartBecauseTooSmall(t *testing.T) {
 	store := New("bucket", s3obj)
 
 	gomock.InOrder(
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.info"),
 		}).Return(&s3.GetObjectOutput{
 			Body: ioutil.NopCloser(bytes.NewReader([]byte(`{"ID":"uploadId","Size":500,"Offset":0,"MetaData":null,"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":null}`))),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -702,11 +702,11 @@ func TestWriteChunkWriteIncompletePartBecauseTooSmall(t *testing.T) {
 				},
 			},
 		}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{}, awserr.New("NoSuchKey", "The specified key does not exist", nil)),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -721,11 +721,11 @@ func TestWriteChunkWriteIncompletePartBecauseTooSmall(t *testing.T) {
 				},
 			},
 		}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{}, awserr.New("NoSuchKey", "The specified key does not exist.", nil)),
-		s3obj.EXPECT().PutObject(NewPutObjectInputMatcher(&s3.PutObjectInput{
+		s3obj.EXPECT().PutObjectWithContext(context.Background(), NewPutObjectInputMatcher(&s3.PutObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 			Body:   bytes.NewReader([]byte("1234567890")),
@@ -753,50 +753,50 @@ func TestWriteChunkPrependsIncompletePart(t *testing.T) {
 	store.MaxObjectSize = 5 * 1024 * 1024 * 1024 * 1024
 
 	gomock.InOrder(
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.info"),
 		}).Return(&s3.GetObjectOutput{
 			Body: ioutil.NopCloser(bytes.NewReader([]byte(`{"ID":"uploadId","Size":5,"Offset":0,"MetaData":null,"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":null}`))),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
 			PartNumberMarker: aws.Int64(0),
 		}).Return(&s3.ListPartsOutput{Parts: []*s3.Part{}}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{
 			ContentLength: aws.Int64(3),
 			Body:          ioutil.NopCloser(bytes.NewReader([]byte("123"))),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
 			PartNumberMarker: aws.Int64(0),
 		}).Return(&s3.ListPartsOutput{Parts: []*s3.Part{}}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{
 			ContentLength: aws.Int64(3),
 			Body:          ioutil.NopCloser(bytes.NewReader([]byte("123"))),
 		}, nil),
-		s3obj.EXPECT().DeleteObject(&s3.DeleteObjectInput{
+		s3obj.EXPECT().DeleteObjectWithContext(context.Background(), &s3.DeleteObjectInput{
 			Bucket: aws.String(store.Bucket),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.DeleteObjectOutput{}, nil),
-		s3obj.EXPECT().UploadPart(NewUploadPartInputMatcher(&s3.UploadPartInput{
+		s3obj.EXPECT().UploadPartWithContext(context.Background(), NewUploadPartInputMatcher(&s3.UploadPartInput{
 			Bucket:     aws.String("bucket"),
 			Key:        aws.String("uploadId"),
 			UploadId:   aws.String("multipartId"),
 			PartNumber: aws.Int64(1),
 			Body:       bytes.NewReader([]byte("1234")),
 		})).Return(nil, nil),
-		s3obj.EXPECT().UploadPart(NewUploadPartInputMatcher(&s3.UploadPartInput{
+		s3obj.EXPECT().UploadPartWithContext(context.Background(), NewUploadPartInputMatcher(&s3.UploadPartInput{
 			Bucket:     aws.String("bucket"),
 			Key:        aws.String("uploadId"),
 			UploadId:   aws.String("multipartId"),
@@ -826,50 +826,50 @@ func TestWriteChunkPrependsIncompletePartAndWritesANewIncompletePart(t *testing.
 	store.MaxObjectSize = 5 * 1024 * 1024 * 1024 * 1024
 
 	gomock.InOrder(
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.info"),
 		}).Return(&s3.GetObjectOutput{
 			Body: ioutil.NopCloser(bytes.NewReader([]byte(`{"ID":"uploadId","Size":10,"Offset":0,"MetaData":null,"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":null}`))),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
 			PartNumberMarker: aws.Int64(0),
 		}).Return(&s3.ListPartsOutput{Parts: []*s3.Part{}}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{
 			ContentLength: aws.Int64(3),
 			Body:          ioutil.NopCloser(bytes.NewReader([]byte("123"))),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
 			PartNumberMarker: aws.Int64(0),
 		}).Return(&s3.ListPartsOutput{Parts: []*s3.Part{}}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{
 			Body:          ioutil.NopCloser(bytes.NewReader([]byte("123"))),
 			ContentLength: aws.Int64(3),
 		}, nil),
-		s3obj.EXPECT().DeleteObject(&s3.DeleteObjectInput{
+		s3obj.EXPECT().DeleteObjectWithContext(context.Background(), &s3.DeleteObjectInput{
 			Bucket: aws.String(store.Bucket),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.DeleteObjectOutput{}, nil),
-		s3obj.EXPECT().UploadPart(NewUploadPartInputMatcher(&s3.UploadPartInput{
+		s3obj.EXPECT().UploadPartWithContext(context.Background(), NewUploadPartInputMatcher(&s3.UploadPartInput{
 			Bucket:     aws.String("bucket"),
 			Key:        aws.String("uploadId"),
 			UploadId:   aws.String("multipartId"),
 			PartNumber: aws.Int64(1),
 			Body:       bytes.NewReader([]byte("1234")),
 		})).Return(nil, nil),
-		s3obj.EXPECT().PutObject(NewPutObjectInputMatcher(&s3.PutObjectInput{
+		s3obj.EXPECT().PutObjectWithContext(context.Background(), NewPutObjectInputMatcher(&s3.PutObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 			Body:   bytes.NewReader([]byte("5")),
@@ -894,13 +894,13 @@ func TestWriteChunkAllowTooSmallLast(t *testing.T) {
 	store.MinPartSize = 20
 
 	gomock.InOrder(
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.info"),
 		}).Return(&s3.GetObjectOutput{
 			Body: ioutil.NopCloser(bytes.NewReader([]byte(`{"ID":"uploadId","Size":500,"Offset":0,"MetaData":null,"IsPartial":false,"IsFinal":false,"PartialUploads":null,"Storage":null}`))),
 		}, nil),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -915,11 +915,11 @@ func TestWriteChunkAllowTooSmallLast(t *testing.T) {
 				},
 			},
 		}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{}, awserr.New("AccessDenied", "Access Denied.", nil)),
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -934,11 +934,11 @@ func TestWriteChunkAllowTooSmallLast(t *testing.T) {
 				},
 			},
 		}, nil),
-		s3obj.EXPECT().GetObject(&s3.GetObjectInput{
+		s3obj.EXPECT().GetObjectWithContext(context.Background(), &s3.GetObjectInput{
 			Bucket: aws.String("bucket"),
 			Key:    aws.String("uploadId.part"),
 		}).Return(&s3.GetObjectOutput{}, awserr.New("NoSuchKey", "The specified key does not exist.", nil)),
-		s3obj.EXPECT().UploadPart(NewUploadPartInputMatcher(&s3.UploadPartInput{
+		s3obj.EXPECT().UploadPartWithContext(context.Background(), NewUploadPartInputMatcher(&s3.UploadPartInput{
 			Bucket:     aws.String("bucket"),
 			Key:        aws.String("uploadId"),
 			UploadId:   aws.String("multipartId"),
@@ -967,13 +967,13 @@ func TestTerminate(t *testing.T) {
 	store := New("bucket", s3obj)
 
 	// Order is not important in this situation.
-	s3obj.EXPECT().AbortMultipartUpload(&s3.AbortMultipartUploadInput{
+	s3obj.EXPECT().AbortMultipartUploadWithContext(context.Background(), &s3.AbortMultipartUploadInput{
 		Bucket:   aws.String("bucket"),
 		Key:      aws.String("uploadId"),
 		UploadId: aws.String("multipartId"),
 	}).Return(nil, nil)
 
-	s3obj.EXPECT().DeleteObjects(&s3.DeleteObjectsInput{
+	s3obj.EXPECT().DeleteObjectsWithContext(context.Background(), &s3.DeleteObjectsInput{
 		Bucket: aws.String("bucket"),
 		Delete: &s3.Delete{
 			Objects: []*s3.ObjectIdentifier{
@@ -1008,13 +1008,13 @@ func TestTerminateWithErrors(t *testing.T) {
 
 	// Order is not important in this situation.
 	// NoSuchUpload errors should be ignored
-	s3obj.EXPECT().AbortMultipartUpload(&s3.AbortMultipartUploadInput{
+	s3obj.EXPECT().AbortMultipartUploadWithContext(context.Background(), &s3.AbortMultipartUploadInput{
 		Bucket:   aws.String("bucket"),
 		Key:      aws.String("uploadId"),
 		UploadId: aws.String("multipartId"),
 	}).Return(nil, awserr.New("NoSuchUpload", "The specified upload does not exist.", nil))
 
-	s3obj.EXPECT().DeleteObjects(&s3.DeleteObjectsInput{
+	s3obj.EXPECT().DeleteObjectsWithContext(context.Background(), &s3.DeleteObjectsInput{
 		Bucket: aws.String("bucket"),
 		Delete: &s3.Delete{
 			Objects: []*s3.ObjectIdentifier{
@@ -1055,7 +1055,7 @@ func TestConcatUploads(t *testing.T) {
 	s3obj := NewMockS3API(mockCtrl)
 	store := New("bucket", s3obj)
 
-	s3obj.EXPECT().UploadPartCopy(&s3.UploadPartCopyInput{
+	s3obj.EXPECT().UploadPartCopyWithContext(context.Background(), &s3.UploadPartCopyInput{
 		Bucket:     aws.String("bucket"),
 		Key:        aws.String("uploadId"),
 		UploadId:   aws.String("multipartId"),
@@ -1063,7 +1063,7 @@ func TestConcatUploads(t *testing.T) {
 		PartNumber: aws.Int64(1),
 	}).Return(nil, nil)
 
-	s3obj.EXPECT().UploadPartCopy(&s3.UploadPartCopyInput{
+	s3obj.EXPECT().UploadPartCopyWithContext(context.Background(), &s3.UploadPartCopyInput{
 		Bucket:     aws.String("bucket"),
 		Key:        aws.String("uploadId"),
 		UploadId:   aws.String("multipartId"),
@@ -1071,7 +1071,7 @@ func TestConcatUploads(t *testing.T) {
 		PartNumber: aws.Int64(2),
 	}).Return(nil, nil)
 
-	s3obj.EXPECT().UploadPartCopy(&s3.UploadPartCopyInput{
+	s3obj.EXPECT().UploadPartCopyWithContext(context.Background(), &s3.UploadPartCopyInput{
 		Bucket:     aws.String("bucket"),
 		Key:        aws.String("uploadId"),
 		UploadId:   aws.String("multipartId"),
@@ -1081,7 +1081,7 @@ func TestConcatUploads(t *testing.T) {
 
 	// Output from s3Store.FinishUpload
 	gomock.InOrder(
-		s3obj.EXPECT().ListParts(&s3.ListPartsInput{
+		s3obj.EXPECT().ListPartsWithContext(context.Background(), &s3.ListPartsInput{
 			Bucket:           aws.String("bucket"),
 			Key:              aws.String("uploadId"),
 			UploadId:         aws.String("multipartId"),
@@ -1102,7 +1102,7 @@ func TestConcatUploads(t *testing.T) {
 				},
 			},
 		}, nil),
-		s3obj.EXPECT().CompleteMultipartUpload(&s3.CompleteMultipartUploadInput{
+		s3obj.EXPECT().CompleteMultipartUploadWithContext(context.Background(), &s3.CompleteMultipartUploadInput{
 			Bucket:   aws.String("bucket"),
 			Key:      aws.String("uploadId"),
 			UploadId: aws.String("multipartId"),