Add tests for S3Store

This commit is contained in:
Marius 2016-01-05 18:21:53 +01:00
parent b6d67debee
commit 2073521776
4 changed files with 1821 additions and 37 deletions

View File

@ -10,7 +10,6 @@ import (
"bytes"
"encoding/json"
"errors"
"fmt"
"io"
"io/ioutil"
"os"
@ -21,34 +20,39 @@ import (
"github.com/aws/aws-sdk-go/aws"
"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/aws/aws-sdk-go/service/s3/s3iface"
)
const MinPartSize = int64(5 * 1024 * 1024)
// See the tusd.DataStore interface for documentation about the different
// methods.
type S3Store struct {
Bucket string
Service *s3.S3
Service s3iface.S3API
MaxPartSize int64
MinPartSize int64
}
func New(bucket string, service *s3.S3) *S3Store {
func New(bucket string, service s3iface.S3API) *S3Store {
return &S3Store{
Bucket: bucket,
Service: service,
MaxPartSize: 6 * 1024 * 1024,
MinPartSize: 5 * 1024 * 1024,
}
}
func (store S3Store) NewUpload(info tusd.FileInfo) (id string, err error) {
uploadId := uid.Uid()
var uploadId string
if info.ID == "" {
uploadId = uid.Uid()
} else {
uploadId = info.ID
}
infoJson, err := json.Marshal(info)
if err != nil {
return id, err
return "", err
}
// Create object on S3 containing information about the file
@ -59,7 +63,7 @@ func (store S3Store) NewUpload(info tusd.FileInfo) (id string, err error) {
ContentLength: aws.Int64(int64(len(infoJson))),
})
if err != nil {
return id, err
return "", err
}
// Create the actual multipart upload
@ -68,7 +72,7 @@ func (store S3Store) NewUpload(info tusd.FileInfo) (id string, err error) {
Key: aws.String(uploadId),
})
if err != nil {
return id, err
return "", err
}
id = uploadId + "+" + *res.UploadId
@ -88,9 +92,6 @@ func (store S3Store) WriteChunk(id string, offset int64, src io.Reader) (int64,
size := info.Size
bytesUploaded := int64(0)
fmt.Println("size:", size)
fmt.Println("offset:", offset)
// Get number of parts to generate next number
listPtr, err := store.Service.ListParts(&s3.ListPartsInput{
Bucket: aws.String(store.Bucket),
@ -103,11 +104,8 @@ func (store S3Store) WriteChunk(id string, offset int64, src io.Reader) (int64,
list := *listPtr
numParts := len(list.Parts)
fmt.Println("numParts:", numParts)
nextPartNum := int64(numParts + 1)
fmt.Println("nextNum:", nextPartNum)
for {
// Create a temporary file to store the part in it
file, err := ioutil.TempFile("", "tusd-s3-tmp-")
@ -117,28 +115,20 @@ func (store S3Store) WriteChunk(id string, offset int64, src io.Reader) (int64,
defer os.Remove(file.Name())
defer file.Close()
fmt.Println("max:", store.MaxPartSize)
limitedReader := io.LimitReader(src, store.MaxPartSize)
n, err := io.Copy(file, limitedReader)
fmt.Println("err:", err)
if err != nil && err != io.EOF {
return bytesUploaded, err
}
fmt.Println("n:", n)
fmt.Println("offset:", offset)
if (size - offset) <= MinPartSize {
if (size - offset) <= store.MinPartSize {
if (size - offset) != n {
return bytesUploaded, nil
}
} else if n < MinPartSize {
} else if n < store.MinPartSize {
return bytesUploaded, nil
}
fmt.Println("upload…")
// Seek to the beginning of the file
file.Seek(0, 0)
@ -153,17 +143,10 @@ func (store S3Store) WriteChunk(id string, offset int64, src io.Reader) (int64,
return bytesUploaded, err
}
pos, _ := file.Seek(1, 0)
fmt.Println("pos:", pos)
offset += bytesUploaded
bytesUploaded += n
nextPartNum += 1
}
fmt.Println("bytes:", bytesUploaded)
return bytesUploaded, nil
}
func (store S3Store) GetInfo(id string) (info tusd.FileInfo, err error) {
@ -197,7 +180,6 @@ func (store S3Store) GetInfo(id string) (info tusd.FileInfo, err error) {
// when the multipart upload has already been completed or aborted. Since
// we already found the info object, we know that the upload has been
// completed and therefore can ensure the the offset is the size.
fmt.Println(err, err.(awserr.Error).Code())
if err, ok := err.(awserr.Error); ok && err.Code() == "NoSuchUpload" {
info.Offset = info.Size
return info, nil
@ -238,7 +220,7 @@ func (store S3Store) GetReader(id string) (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.ListParts(&s3.ListPartsInput{
Bucket: aws.String(store.Bucket),
Key: aws.String(uploadId),
UploadId: aws.String(multipartId),
@ -284,8 +266,6 @@ func (store S3Store) Terminate(id string) error {
func (store S3Store) FinishUpload(id string) error {
uploadId, multipartId := splitIds(id)
println("Finish upload")
// Get uploaded parts
listPtr, err := store.Service.ListParts(&s3.ListPartsInput{
Bucket: aws.String(store.Bucket),

1303
s3store/s3store_mock_test.go Normal file

File diff suppressed because it is too large Load Diff

444
s3store/s3store_test.go Normal file
View File

@ -0,0 +1,444 @@
package s3store_test
import (
"bytes"
"io/ioutil"
"testing"
"github.com/golang/mock/gomock"
"github.com/stretchr/testify/assert"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/awserr"
"github.com/aws/aws-sdk-go/service/s3"
"github.com/tus/tusd"
"github.com/tus/tusd/s3store"
)
//go:generate mockgen -destination=./s3store_mock_test.go -package=s3store_test github.com/aws/aws-sdk-go/service/s3/s3iface S3API
func TestNewUpload(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)
s3obj := NewMockS3API(mockCtrl)
store := s3store.New("bucket", s3obj)
assert.Equal(store.Bucket, "bucket")
assert.Equal(store.Service, s3obj)
gomock.InOrder(
s3obj.EXPECT().PutObject(&s3.PutObjectInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId.info"),
Body: bytes.NewReader([]byte(`{"ID":"uploadId","Size":500,"Offset":0,"MetaData":null,"IsPartial":false,"IsFinal":false,"PartialUploads":null}`)),
ContentLength: aws.Int64(int64(111)),
}),
s3obj.EXPECT().CreateMultipartUpload(&s3.CreateMultipartUploadInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
}).Return(&s3.CreateMultipartUploadOutput{
UploadId: aws.String("multipartId"),
}, nil),
)
info := tusd.FileInfo{
ID: "uploadId",
Size: 500,
}
id, err := store.NewUpload(info)
assert.Nil(err)
assert.Equal(id, "uploadId+multipartId")
}
func TestGetInfoNotFound(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)
s3obj := NewMockS3API(mockCtrl)
store := s3store.New("bucket", s3obj)
s3obj.EXPECT().GetObject(&s3.GetObjectInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId.info"),
}).Return(nil, awserr.New("NoSuchKey", "The specified key does not exist.", nil))
_, err := store.GetInfo("uploadId+multipartId")
assert.Equal(err, tusd.ErrNotFound)
}
func TestGetInfo(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)
s3obj := NewMockS3API(mockCtrl)
store := s3store.New("bucket", s3obj)
gomock.InOrder(
s3obj.EXPECT().GetObject(&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}`))),
}, nil),
s3obj.EXPECT().ListParts(&s3.ListPartsInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
}).Return(&s3.ListPartsOutput{
Parts: []*s3.Part{
{
Size: aws.Int64(100),
},
{
Size: aws.Int64(200),
},
},
}, nil),
)
info, err := store.GetInfo("uploadId+multipartId")
assert.Nil(err)
assert.Equal(int64(500), info.Size)
assert.Equal(int64(300), info.Offset)
}
func TestGetInfoFinished(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)
s3obj := NewMockS3API(mockCtrl)
store := s3store.New("bucket", s3obj)
gomock.InOrder(
s3obj.EXPECT().GetObject(&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}`))),
}, nil),
s3obj.EXPECT().ListParts(&s3.ListPartsInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
}).Return(nil, awserr.New("NoSuchUpload", "The specified upload does not exist.", nil)),
)
info, err := store.GetInfo("uploadId+multipartId")
assert.Nil(err)
assert.Equal(int64(500), info.Size)
assert.Equal(int64(500), info.Offset)
}
func TestGetReader(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)
s3obj := NewMockS3API(mockCtrl)
store := s3store.New("bucket", s3obj)
s3obj.EXPECT().GetObject(&s3.GetObjectInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
}).Return(&s3.GetObjectOutput{
Body: ioutil.NopCloser(bytes.NewReader([]byte(`hello world`))),
}, nil)
content, err := store.GetReader("uploadId+multipartId")
assert.Nil(err)
assert.Equal(content, ioutil.NopCloser(bytes.NewReader([]byte(`hello world`))))
}
func TestGetReaderNotFound(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)
s3obj := NewMockS3API(mockCtrl)
store := s3store.New("bucket", s3obj)
gomock.InOrder(
s3obj.EXPECT().GetObject(&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{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
MaxParts: aws.Int64(0),
}).Return(nil, awserr.New("NoSuchUpload", "The specified upload does not exist.", nil)),
)
content, err := store.GetReader("uploadId+multipartId")
assert.Nil(content)
assert.Equal(err, tusd.ErrNotFound)
}
func TestGetReaderNotFinished(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)
s3obj := NewMockS3API(mockCtrl)
store := s3store.New("bucket", s3obj)
gomock.InOrder(
s3obj.EXPECT().GetObject(&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{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
MaxParts: aws.Int64(0),
}).Return(&s3.ListPartsOutput{
Parts: []*s3.Part{},
}, nil),
)
content, err := store.GetReader("uploadId+multipartId")
assert.Nil(content)
assert.Equal(err.Error(), "cannot stream non-finished upload")
}
func TestFinishUpload(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)
s3obj := NewMockS3API(mockCtrl)
store := s3store.New("bucket", s3obj)
gomock.InOrder(
s3obj.EXPECT().ListParts(&s3.ListPartsInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
}).Return(&s3.ListPartsOutput{
Parts: []*s3.Part{
{
Size: aws.Int64(100),
ETag: aws.String("foo"),
PartNumber: aws.Int64(1),
},
{
Size: aws.Int64(200),
ETag: aws.String("bar"),
PartNumber: aws.Int64(2),
},
},
}, nil),
s3obj.EXPECT().CompleteMultipartUpload(&s3.CompleteMultipartUploadInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
MultipartUpload: &s3.CompletedMultipartUpload{
Parts: []*s3.CompletedPart{
{
ETag: aws.String("foo"),
PartNumber: aws.Int64(1),
},
{
ETag: aws.String("bar"),
PartNumber: aws.Int64(2),
},
},
},
}).Return(nil, nil),
)
err := store.FinishUpload("uploadId+multipartId")
assert.Nil(err)
}
func TestWriteChunk(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)
s3obj := NewMockS3API(mockCtrl)
store := s3store.New("bucket", s3obj)
store.MaxPartSize = 4
store.MinPartSize = 2
gomock.InOrder(
s3obj.EXPECT().GetObject(&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}`))),
}, nil),
s3obj.EXPECT().ListParts(&s3.ListPartsInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
}).Return(&s3.ListPartsOutput{
Parts: []*s3.Part{
{
Size: aws.Int64(100),
},
{
Size: aws.Int64(200),
},
},
}, nil),
s3obj.EXPECT().ListParts(&s3.ListPartsInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
}).Return(&s3.ListPartsOutput{
Parts: []*s3.Part{
{
Size: aws.Int64(100),
},
{
Size: aws.Int64(200),
},
},
}, nil),
s3obj.EXPECT().UploadPart(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{
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{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
PartNumber: aws.Int64(5),
Body: bytes.NewReader([]byte("90")),
})).Return(nil, nil),
)
bytesRead, err := store.WriteChunk("uploadId+multipartId", 300, bytes.NewReader([]byte("1234567890")))
assert.Nil(err)
assert.Equal(int64(10), bytesRead)
}
func TestWriteChunkDropTooSmall(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)
s3obj := NewMockS3API(mockCtrl)
store := s3store.New("bucket", s3obj)
gomock.InOrder(
s3obj.EXPECT().GetObject(&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}`))),
}, nil),
s3obj.EXPECT().ListParts(&s3.ListPartsInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
}).Return(&s3.ListPartsOutput{
Parts: []*s3.Part{
{
Size: aws.Int64(100),
},
{
Size: aws.Int64(200),
},
},
}, nil),
s3obj.EXPECT().ListParts(&s3.ListPartsInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
}).Return(&s3.ListPartsOutput{
Parts: []*s3.Part{
{
Size: aws.Int64(100),
},
{
Size: aws.Int64(200),
},
},
}, nil),
)
bytesRead, err := store.WriteChunk("uploadId+multipartId", 300, bytes.NewReader([]byte("1234567890")))
assert.Nil(err)
assert.Equal(int64(0), bytesRead)
}
func TestWriteChunkAllowTooSmallLast(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)
s3obj := NewMockS3API(mockCtrl)
store := s3store.New("bucket", s3obj)
store.MinPartSize = 20
gomock.InOrder(
s3obj.EXPECT().GetObject(&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}`))),
}, nil),
s3obj.EXPECT().ListParts(&s3.ListPartsInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
}).Return(&s3.ListPartsOutput{
Parts: []*s3.Part{
{
Size: aws.Int64(400),
},
{
Size: aws.Int64(90),
},
},
}, nil),
s3obj.EXPECT().ListParts(&s3.ListPartsInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
}).Return(&s3.ListPartsOutput{
Parts: []*s3.Part{
{
Size: aws.Int64(400),
},
{
Size: aws.Int64(90),
},
},
}, nil),
s3obj.EXPECT().UploadPart(NewUploadPartInputMatcher(&s3.UploadPartInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
UploadId: aws.String("multipartId"),
PartNumber: aws.Int64(3),
Body: bytes.NewReader([]byte("1234567890")),
})).Return(nil, nil),
)
// 10 bytes are missing for the upload to be finished (offset at 490 for 500
// bytes file) but the minimum chunk size is higher (20). The chunk is
// still uploaded since the last part may be smaller than the minimum.
bytesRead, err := store.WriteChunk("uploadId+multipartId", 490, bytes.NewReader([]byte("1234567890")))
assert.Nil(err)
assert.Equal(int64(10), bytesRead)
}

View File

@ -0,0 +1,57 @@
package s3store_test
import (
"fmt"
"io/ioutil"
"reflect"
"github.com/aws/aws-sdk-go/service/s3"
"github.com/golang/mock/gomock"
)
type UploadPartInputMatcher struct {
expect *s3.UploadPartInput
}
func NewUploadPartInputMatcher(expect *s3.UploadPartInput) gomock.Matcher {
return UploadPartInputMatcher{
expect: expect,
}
}
func (m UploadPartInputMatcher) Matches(x interface{}) bool {
input, ok := x.(*s3.UploadPartInput)
if !ok {
return false
}
inputBody := input.Body
expectBody := m.expect.Body
i, err := ioutil.ReadAll(inputBody)
if err != nil {
panic(err)
}
inputBody.Seek(0, 0)
e, err := ioutil.ReadAll(expectBody)
if err != nil {
panic(err)
}
m.expect.Body.Seek(0, 0)
if !reflect.DeepEqual(e, i) {
return false
}
input.Body = nil
m.expect.Body = nil
return reflect.DeepEqual(m.expect, input)
}
func (m UploadPartInputMatcher) String() string {
body, _ := ioutil.ReadAll(m.expect.Body)
m.expect.Body.Seek(0, 0)
return fmt.Sprintf("UploadPartInput(%d: %s)", *m.expect.PartNumber, body)
}