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 }