github.com/lyft/flytestdlib@v0.3.12-0.20210213045714-8cdd111ecda1/storage/protobuf_store_test.go (about)

     1  package storage
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"math/rand"
     8  	"testing"
     9  
    10  	"github.com/lyft/flytestdlib/promutils"
    11  
    12  	"github.com/golang/protobuf/proto"
    13  	errs "github.com/pkg/errors"
    14  	"github.com/stretchr/testify/assert"
    15  )
    16  
    17  type mockProtoMessage struct {
    18  	X int64 `protobuf:"varint,2,opt,name=x,json=x,proto3" json:"x,omitempty"`
    19  }
    20  
    21  type mockBigDataProtoMessage struct {
    22  	X []byte `protobuf:"bytes,1,opt,name=X,proto3" json:"X,omitempty"`
    23  }
    24  
    25  func (mockProtoMessage) Reset() {
    26  }
    27  
    28  func (m mockProtoMessage) String() string {
    29  	return proto.CompactTextString(m)
    30  }
    31  
    32  func (mockProtoMessage) ProtoMessage() {
    33  }
    34  
    35  func (mockBigDataProtoMessage) Reset() {
    36  }
    37  
    38  func (m mockBigDataProtoMessage) String() string {
    39  	return proto.CompactTextString(m)
    40  }
    41  
    42  func (mockBigDataProtoMessage) ProtoMessage() {
    43  }
    44  
    45  func TestDefaultProtobufStore_ReadProtobuf(t *testing.T) {
    46  	t.Run("Read after Write", func(t *testing.T) {
    47  		testScope := promutils.NewTestScope()
    48  		s, err := NewDataStore(&Config{Type: TypeMemory}, testScope)
    49  		assert.NoError(t, err)
    50  
    51  		err = s.WriteProtobuf(context.TODO(), DataReference("hello"), Options{}, &mockProtoMessage{X: 5})
    52  		assert.NoError(t, err)
    53  
    54  		m := &mockProtoMessage{}
    55  		err = s.ReadProtobuf(context.TODO(), DataReference("hello"), m)
    56  		assert.NoError(t, err)
    57  		assert.Equal(t, int64(5), m.X)
    58  	})
    59  }
    60  
    61  func TestDefaultProtobufStore_BigDataReadAfterWrite(t *testing.T) {
    62  	t.Run("Read after Write with Big Data", func(t *testing.T) {
    63  		testScope := promutils.NewTestScope()
    64  
    65  		s, err := NewDataStore(
    66  			&Config{
    67  				Type: TypeMemory,
    68  				Cache: CachingConfig{
    69  					MaxSizeMegabytes: 1,
    70  					TargetGCPercent:  20,
    71  				},
    72  			}, testScope)
    73  		assert.NoError(t, err)
    74  
    75  		bigD := make([]byte, 1.5*1024*1024)
    76  		// #nosec G404
    77  		_, err = rand.Read(bigD)
    78  		assert.NoError(t, err)
    79  
    80  		mockMessage := &mockBigDataProtoMessage{X: bigD}
    81  
    82  		err = s.WriteProtobuf(context.TODO(), DataReference("bigK"), Options{}, mockMessage)
    83  		assert.NoError(t, err)
    84  
    85  		m := &mockBigDataProtoMessage{}
    86  		err = s.ReadProtobuf(context.TODO(), DataReference("bigK"), m)
    87  		assert.NoError(t, err)
    88  		assert.Equal(t, bigD, m.X)
    89  
    90  	})
    91  }
    92  
    93  func TestDefaultProtobufStore_HardErrors(t *testing.T) {
    94  	ctx := context.TODO()
    95  	k1 := DataReference("k1")
    96  	dummyHeadErrorMsg := "Dummy head error"
    97  	dummyWriteErrorMsg := "Dummy write error"
    98  	dummyReadErrorMsg := "Dummy read error"
    99  	store := &dummyStore{
   100  		HeadCb: func(ctx context.Context, reference DataReference) (Metadata, error) {
   101  			return MemoryMetadata{}, fmt.Errorf(dummyHeadErrorMsg)
   102  		},
   103  		WriteRawCb: func(ctx context.Context, reference DataReference, size int64, opts Options, raw io.Reader) error {
   104  			return fmt.Errorf(dummyWriteErrorMsg)
   105  		},
   106  		ReadRawCb: func(ctx context.Context, reference DataReference) (io.ReadCloser, error) {
   107  			return nil, fmt.Errorf(dummyReadErrorMsg)
   108  		},
   109  	}
   110  	testScope := promutils.NewTestScope()
   111  	pbErroneousStore := NewDefaultProtobufStore(store, testScope)
   112  	t.Run("Test if hard write errors are handled correctly", func(t *testing.T) {
   113  		err := pbErroneousStore.WriteProtobuf(ctx, k1, Options{}, &mockProtoMessage{X: 5})
   114  		assert.False(t, IsFailedWriteToCache(err))
   115  		assert.Equal(t, dummyWriteErrorMsg, errs.Cause(err).Error())
   116  	})
   117  
   118  	t.Run("Test if hard read errors are handled correctly", func(t *testing.T) {
   119  		m := &mockProtoMessage{}
   120  		err := pbErroneousStore.ReadProtobuf(ctx, k1, m)
   121  		assert.False(t, IsFailedWriteToCache(err))
   122  		assert.Equal(t, dummyReadErrorMsg, errs.Cause(err).Error())
   123  	})
   124  }