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