github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/graveler/committed/batch_test.go (about) 1 package committed_test 2 3 import ( 4 "errors" 5 "fmt" 6 "strconv" 7 "testing" 8 9 "github.com/go-test/deep" 10 "github.com/stretchr/testify/assert" 11 "github.com/stretchr/testify/require" 12 "github.com/treeverse/lakefs/pkg/graveler" 13 "github.com/treeverse/lakefs/pkg/graveler/committed" 14 ) 15 16 // FakeRangeWriter is a RangeWriter that is safe to use in goroutines. (mock.RangeWriter is 17 // NOT safe, it uses gomock which calls t.Fatal and friends!) 18 type FakeRangeWriter struct { 19 err error 20 21 closed bool 22 closeResult closeResult 23 writeRecords []*committed.Record 24 25 storedType string 26 } 27 28 type closeResult struct { 29 result *committed.WriteResult 30 err error 31 } 32 33 var ( 34 ErrAlreadyClosed = errors.New("closed more than once") 35 ErrUnexpected = errors.New("unexpected call") 36 ) 37 38 func (f *FakeRangeWriter) Err() error { 39 return f.err 40 } 41 42 func (f *FakeRangeWriter) setErr(err error) error { 43 f.err = err 44 return err 45 } 46 47 func (f *FakeRangeWriter) ExpectWriteRecord(r committed.Record) { 48 f.writeRecords = append(f.writeRecords, &r) 49 } 50 51 func (f *FakeRangeWriter) ExpectAnyRecord() { 52 f.writeRecords = append(f.writeRecords, nil) 53 } 54 55 func (f *FakeRangeWriter) WriteRecord(r committed.Record) error { 56 if len(f.writeRecords) < 1 { 57 return f.setErr(fmt.Errorf("try to write %+v when expected nothing: %w", r, ErrUnexpected)) 58 } 59 var n *committed.Record 60 n, f.writeRecords = f.writeRecords[0], f.writeRecords[1:] 61 if n == nil { // any 62 return nil 63 } 64 if diffs := deep.Equal(&r, n); diffs != nil { 65 return f.setErr(fmt.Errorf("try to write the wrong value %s: %w", diffs, ErrUnexpected)) 66 } 67 return nil 68 } 69 70 func (f *FakeRangeWriter) SetMetadata(key, value string) { 71 if key == committed.MetadataTypeKey { 72 f.storedType = value 73 } 74 } 75 76 func (*FakeRangeWriter) GetApproximateSize() uint64 { return 0 } 77 78 func (*FakeRangeWriter) ShouldBreakAtKey(graveler.Key, *committed.Params) bool { return false } 79 80 func (f *FakeRangeWriter) Close() (*committed.WriteResult, error) { 81 if f.closed { 82 f.err = ErrAlreadyClosed 83 return nil, ErrAlreadyClosed 84 } 85 return f.closeResult.result, f.closeResult.err 86 } 87 88 func (f *FakeRangeWriter) Abort() error { return nil } 89 90 func NewFakeRangeWriter(result *committed.WriteResult, err error) *FakeRangeWriter { 91 return &FakeRangeWriter{ 92 closeResult: closeResult{result, err}, 93 } 94 } 95 96 func TestBatchCloserSuccess(t *testing.T) { 97 runSuccessScenario(t, 502, 5) 98 } 99 100 func TestBatchWriterFailed(t *testing.T) { 101 writerSuccess := NewFakeRangeWriter( 102 &committed.WriteResult{ 103 RangeID: committed.ID(strconv.Itoa(1)), 104 First: committed.Key("row_1"), 105 Last: committed.Key("row_2"), 106 Count: 4321, 107 }, nil) 108 expectedErr := errors.New("failure") 109 writerFailure := NewFakeRangeWriter(nil, expectedErr) 110 111 sut := committed.NewBatchCloser(10) 112 assert.NoError(t, sut.CloseWriterAsync(writerSuccess)) 113 assert.NoError(t, sut.CloseWriterAsync(writerFailure)) 114 115 res, err := sut.Wait() 116 assert.Error(t, expectedErr, err) 117 assert.Nil(t, res) 118 119 assert.NoError(t, writerSuccess.Err()) 120 assert.NoError(t, writerFailure.Err()) 121 } 122 123 func TestBatchCloserMultipleWaitCalls(t *testing.T) { 124 writer := NewFakeRangeWriter(&committed.WriteResult{ 125 RangeID: "last", 126 First: committed.Key("row_1"), 127 Last: committed.Key("row_2"), 128 Count: 4321, 129 }, nil) 130 131 sut := runSuccessScenario(t, 1, 1) 132 133 assert.Error(t, sut.CloseWriterAsync(writer), committed.ErrMultipleWaitCalls) 134 res, err := sut.Wait() 135 require.Nil(t, res) 136 require.Error(t, err, committed.ErrMultipleWaitCalls) 137 } 138 139 func runSuccessScenario(t *testing.T, numWriters, numClosers int) *committed.BatchCloser { 140 t.Helper() 141 142 writers := make([]*FakeRangeWriter, numWriters) 143 for i := 0; i < numWriters; i++ { 144 writers[i] = NewFakeRangeWriter(&committed.WriteResult{ 145 RangeID: committed.ID(strconv.Itoa(i)), 146 First: committed.Key(fmt.Sprintf("row_%d_1", i)), 147 Last: committed.Key(fmt.Sprintf("row_%d_2", i)), 148 Count: i, 149 }, nil) 150 } 151 152 sut := committed.NewBatchCloser(numClosers) 153 154 for i := 0; i < numWriters; i++ { 155 assert.NoError(t, sut.CloseWriterAsync(writers[i])) 156 } 157 158 res, err := sut.Wait() 159 assert.NoError(t, err) 160 assert.NotNil(t, res) 161 assert.Len(t, res, numWriters) 162 163 return sut 164 }