github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/storage/bucket/client_mock.go (about)

     1  package bucket
     2  
     3  import (
     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  
    15  var errObjectDoesNotExist = errors.New("object does not exist")
    16  
    17  // ClientMock mocks objstore.Bucket
    18  type ClientMock struct {
    19  	mock.Mock
    20  }
    21  
    22  // Upload mocks objstore.Bucket.Upload()
    23  func (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  
    28  func (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()
    33  func (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()
    39  func (m *ClientMock) Name() string {
    40  	return "mock"
    41  }
    42  
    43  // Iter mocks objstore.Bucket.Iter()
    44  func (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()
    50  func (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.
    56  func (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()
    73  func (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()
    83  func (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  
   105  func (m *ClientMock) MockDelete(name string, err error) {
   106  	m.On("Delete", mock.Anything, name).Return(err)
   107  }
   108  
   109  func (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()
   114  func (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()
   120  func (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()
   126  func (m *ClientMock) IsObjNotFoundErr(err error) bool {
   127  	return err == errObjectDoesNotExist
   128  }
   129  
   130  // ObjectSize mocks objstore.Bucket.Attributes()
   131  func (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()
   137  func (m *ClientMock) Close() error {
   138  	return nil
   139  }