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  }