github.com/grafana/pyroscope@v1.18.0/pkg/test/mocks/mockobjstore/helper.go (about)

     1  package mockobjstore
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"errors"
     7  	"io"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/stretchr/testify/mock"
    12  	"github.com/thanos-io/objstore"
    13  )
    14  
    15  // ErrObjectDoesNotExist is used in tests to simulate objstore.Bucket.IsObjNotFoundErr().
    16  var ErrObjectDoesNotExist = errors.New("object does not exist")
    17  
    18  func NewMockBucketWithHelper(t testing.TB) *MockBucket {
    19  	m := &MockBucket{}
    20  	m.EXPECT().IsObjNotFoundErr(mock.Anything).RunAndReturn(func(err error) bool {
    21  		return errors.Is(err, ErrObjectDoesNotExist)
    22  	})
    23  	m.EXPECT().Name().Return("mock")
    24  	return m
    25  }
    26  
    27  func (m *MockBucket) MockIter(prefix string, objects []string, err error) {
    28  	m.MockIterWithCallback(prefix, objects, err, nil)
    29  }
    30  
    31  // MockIterWithCallback is a convenient method to mock Iter() and get a callback called when the Iter
    32  // API is called.
    33  func (m *MockBucket) MockIterWithCallback(prefix string, objects []string, err error, cb func()) {
    34  	// implement IsObjNotFoundErr
    35  	m.On("Iter", mock.Anything, prefix, mock.Anything, mock.Anything).Return(err).Run(func(args mock.Arguments) {
    36  		if cb != nil {
    37  			cb()
    38  		}
    39  
    40  		f := args.Get(2).(func(string) error)
    41  
    42  		for _, o := range objects {
    43  			if f(o) != nil {
    44  				break
    45  			}
    46  		}
    47  	})
    48  }
    49  
    50  func (m *MockBucket) MockUpload(name string, err error) {
    51  	m.On("Upload", mock.Anything, name, mock.Anything, mock.Anything).Return(err)
    52  }
    53  
    54  func (m *MockBucket) MockAttributes(name string, attrs objstore.ObjectAttributes, err error) {
    55  	m.On("Attributes", mock.Anything, name).Return(attrs, err)
    56  }
    57  
    58  func (m *MockBucket) MockDelete(name string, err error) {
    59  	m.On("Delete", mock.Anything, name).Return(err)
    60  }
    61  
    62  func (m *MockBucket) MockExists(name string, exists bool, err error) {
    63  	m.On("Exists", mock.Anything, name).Return(exists, err)
    64  }
    65  
    66  // MockGet is a convenient method to mock Get() and Exists()
    67  func (m *MockBucket) MockGet(name, content string, err error) {
    68  	if content != "" {
    69  		m.On("Exists", mock.Anything, name).Return(true, err)
    70  		m.On("Attributes", mock.Anything, name).Return(objstore.ObjectAttributes{
    71  			Size:         int64(len(content)),
    72  			LastModified: time.Now(),
    73  		}, nil)
    74  
    75  		// Since we return an ReadCloser and it can be consumed only once,
    76  		// each time the mocked Get() is called we do create a new one, so
    77  		// that getting the same mocked object twice works as expected.
    78  		m.On("Get", mock.Anything, name).Return(func(_ context.Context, _ string) (io.ReadCloser, error) {
    79  			return io.NopCloser(bytes.NewReader([]byte(content))), err
    80  		})
    81  	} else {
    82  		m.On("Exists", mock.Anything, name).Return(false, err)
    83  		m.On("Get", mock.Anything, name).Return(nil, ErrObjectDoesNotExist)
    84  		m.On("Attributes", mock.Anything, name).Return(nil, ErrObjectDoesNotExist)
    85  	}
    86  }