1package bucket 2 3import ( 4 "bytes" 5 "context" 6 "errors" 7 "io" 8 "io/ioutil" 9 "time" 10 11 "github.com/stretchr/testify/mock" 12 "github.com/thanos-io/thanos/pkg/objstore" 13) 14 15var errObjectDoesNotExist = errors.New("object does not exist") 16 17// ClientMock mocks objstore.Bucket 18type ClientMock struct { 19 mock.Mock 20} 21 22// Upload mocks objstore.Bucket.Upload() 23func (m *ClientMock) Upload(ctx context.Context, name string, r io.Reader) error { 24 args := m.Called(ctx, name, r) 25 return args.Error(0) 26} 27 28func (m *ClientMock) MockUpload(name string, err error) { 29 m.On("Upload", mock.Anything, name, mock.Anything).Return(err) 30} 31 32// Delete mocks objstore.Bucket.Delete() 33func (m *ClientMock) Delete(ctx context.Context, name string) error { 34 args := m.Called(ctx, name) 35 return args.Error(0) 36} 37 38// Name mocks objstore.Bucket.Name() 39func (m *ClientMock) Name() string { 40 return "mock" 41} 42 43// Iter mocks objstore.Bucket.Iter() 44func (m *ClientMock) Iter(ctx context.Context, dir string, f func(string) error, options ...objstore.IterOption) error { 45 args := m.Called(ctx, dir, f, options) 46 return args.Error(0) 47} 48 49// MockIter is a convenient method to mock Iter() 50func (m *ClientMock) MockIter(prefix string, objects []string, err error) { 51 m.MockIterWithCallback(prefix, objects, err, nil) 52} 53 54// MockIterWithCallback is a convenient method to mock Iter() and get a callback called when the Iter 55// API is called. 56func (m *ClientMock) MockIterWithCallback(prefix string, objects []string, err error, cb func()) { 57 m.On("Iter", mock.Anything, prefix, mock.Anything, mock.Anything).Return(err).Run(func(args mock.Arguments) { 58 if cb != nil { 59 cb() 60 } 61 62 f := args.Get(2).(func(string) error) 63 64 for _, o := range objects { 65 if f(o) != nil { 66 break 67 } 68 } 69 }) 70} 71 72// Get mocks objstore.Bucket.Get() 73func (m *ClientMock) Get(ctx context.Context, name string) (io.ReadCloser, error) { 74 args := m.Called(ctx, name) 75 val, err := args.Get(0), args.Error(1) 76 if val == nil { 77 return nil, err 78 } 79 return val.(io.ReadCloser), err 80} 81 82// MockGet is a convenient method to mock Get() and Exists() 83func (m *ClientMock) MockGet(name, content string, err error) { 84 if content != "" { 85 m.On("Exists", mock.Anything, name).Return(true, err) 86 m.On("Attributes", mock.Anything, name).Return(objstore.ObjectAttributes{ 87 Size: int64(len(content)), 88 LastModified: time.Now(), 89 }, nil) 90 91 // Since we return an ReadCloser and it can be consumed only once, 92 // each time the mocked Get() is called we do create a new one, so 93 // that getting the same mocked object twice works as expected. 94 mockedGet := m.On("Get", mock.Anything, name) 95 mockedGet.Run(func(args mock.Arguments) { 96 mockedGet.Return(ioutil.NopCloser(bytes.NewReader([]byte(content))), err) 97 }) 98 } else { 99 m.On("Exists", mock.Anything, name).Return(false, err) 100 m.On("Get", mock.Anything, name).Return(nil, errObjectDoesNotExist) 101 m.On("Attributes", mock.Anything, name).Return(nil, errObjectDoesNotExist) 102 } 103} 104 105func (m *ClientMock) MockDelete(name string, err error) { 106 m.On("Delete", mock.Anything, name).Return(err) 107} 108 109func (m *ClientMock) MockExists(name string, exists bool, err error) { 110 m.On("Exists", mock.Anything, name).Return(exists, err) 111} 112 113// GetRange mocks objstore.Bucket.GetRange() 114func (m *ClientMock) GetRange(ctx context.Context, name string, off, length int64) (io.ReadCloser, error) { 115 args := m.Called(ctx, name, off, length) 116 return args.Get(0).(io.ReadCloser), args.Error(1) 117} 118 119// Exists mocks objstore.Bucket.Exists() 120func (m *ClientMock) Exists(ctx context.Context, name string) (bool, error) { 121 args := m.Called(ctx, name) 122 return args.Bool(0), args.Error(1) 123} 124 125// IsObjNotFoundErr mocks objstore.Bucket.IsObjNotFoundErr() 126func (m *ClientMock) IsObjNotFoundErr(err error) bool { 127 return err == errObjectDoesNotExist 128} 129 130// ObjectSize mocks objstore.Bucket.Attributes() 131func (m *ClientMock) Attributes(ctx context.Context, name string) (objstore.ObjectAttributes, error) { 132 args := m.Called(ctx, name) 133 return args.Get(0).(objstore.ObjectAttributes), args.Error(1) 134} 135 136// Close mocks objstore.Bucket.Close() 137func (m *ClientMock) Close() error { 138 return nil 139} 140