github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/graveler/sstable/range_manager_test.go (about)

     1  package sstable_test
     2  
     3  import (
     4  	"context"
     5  	"crypto"
     6  	"errors"
     7  	"sort"
     8  	"testing"
     9  
    10  	pebblesst "github.com/cockroachdb/pebble/sstable"
    11  
    12  	"github.com/golang/mock/gomock"
    13  	"github.com/stretchr/testify/require"
    14  	"github.com/treeverse/lakefs/pkg/graveler/committed"
    15  	"github.com/treeverse/lakefs/pkg/graveler/sstable"
    16  	fsMock "github.com/treeverse/lakefs/pkg/pyramid/mock"
    17  )
    18  
    19  func makeNewReader(r fakeReader) func(context.Context, committed.Namespace, committed.ID) (*pebblesst.Reader, error) {
    20  	return func(context.Context, committed.Namespace, committed.ID) (*pebblesst.Reader, error) {
    21  		return r.Reader, nil
    22  	}
    23  }
    24  
    25  type NoCache struct{}
    26  
    27  func (n *NoCache) Unref() {}
    28  
    29  func TestGetEntrySuccess(t *testing.T) {
    30  	ctx := context.Background()
    31  	ctrl := gomock.NewController(t)
    32  
    33  	mockFS := fsMock.NewMockFS(ctrl)
    34  
    35  	keys := randomStrings(10)
    36  	sort.Strings(keys)
    37  	vals := randomStrings(len(keys))
    38  
    39  	reader := createSStableReader(t, keys, vals)
    40  
    41  	sut := sstable.NewPebbleSSTableRangeManagerWithNewReader(makeNewReader(reader), &NoCache{}, mockFS, crypto.SHA256)
    42  
    43  	ns := "some-ns"
    44  	sstableID := "some-id"
    45  
    46  	val, err := sut.GetValue(ctx, committed.Namespace(ns), committed.ID(sstableID), committed.Key(keys[len(keys)/3]))
    47  	require.NoError(t, err)
    48  	require.NotNil(t, val)
    49  
    50  	require.Equal(t, 1, reader.GetNumClosed())
    51  }
    52  
    53  func TestGetEntryCacheFailure(t *testing.T) {
    54  	ctx := context.Background()
    55  	ctrl := gomock.NewController(t)
    56  
    57  	expectedErr := errors.New("cache failure")
    58  
    59  	mockFS := fsMock.NewMockFS(ctrl)
    60  
    61  	sut := sstable.NewPebbleSSTableRangeManagerWithNewReader(func(context.Context, committed.Namespace, committed.ID) (*pebblesst.Reader, error) {
    62  		return nil, expectedErr
    63  	}, &NoCache{}, mockFS, crypto.SHA256)
    64  
    65  	ns := "some-ns"
    66  	sstableID := committed.ID("some-id")
    67  
    68  	val, err := sut.GetValue(ctx, committed.Namespace(ns), sstableID, committed.Key("some-key"))
    69  	require.Error(t, expectedErr, err)
    70  	require.Nil(t, val)
    71  }
    72  
    73  func TestGetEntryNotFound(t *testing.T) {
    74  	ctx := context.Background()
    75  	ctrl := gomock.NewController(t)
    76  
    77  	mockFS := fsMock.NewMockFS(ctrl)
    78  
    79  	keys := randomStrings(10)
    80  	sort.Strings(keys)
    81  	vals := randomStrings(len(keys))
    82  
    83  	reader := createSStableReader(t, keys, vals)
    84  
    85  	sut := sstable.NewPebbleSSTableRangeManagerWithNewReader(makeNewReader(reader), &NoCache{}, mockFS, crypto.SHA256)
    86  
    87  	ns := "some-ns"
    88  	sstableID := committed.ID("some-id")
    89  
    90  	val, err := sut.GetValue(ctx, committed.Namespace(ns), sstableID, committed.Key("does-not-exist"))
    91  	require.Error(t, err)
    92  	require.Nil(t, val)
    93  
    94  	require.Equal(t, 1, reader.GetNumClosed())
    95  }
    96  
    97  func TestGetWriterSuccess(t *testing.T) {
    98  	ctx := context.Background()
    99  	ctrl := gomock.NewController(t)
   100  
   101  	mockFS := fsMock.NewMockFS(ctrl)
   102  
   103  	sut := sstable.NewPebbleSSTableRangeManagerWithNewReader(nil, &NoCache{}, mockFS, crypto.SHA256)
   104  
   105  	ns := "some-ns"
   106  	mockFile := fsMock.NewMockStoredFile(ctrl)
   107  	mockFS.EXPECT().Create(ctx, ns).Return(mockFile, nil).Times(1)
   108  
   109  	writer, err := sut.GetWriter(ctx, committed.Namespace(ns), nil)
   110  	require.NoError(t, err)
   111  	require.NotNil(t, writer)
   112  
   113  	require.IsType(t, &sstable.DiskWriter{}, writer)
   114  	dw := writer.(*sstable.DiskWriter)
   115  
   116  	require.Equal(t, mockFS, dw.GetFS())
   117  	require.Equal(t, mockFile, dw.GetStoredFile())
   118  }
   119  
   120  func TestNewPartIteratorSuccess(t *testing.T) {
   121  	ctx := context.Background()
   122  	ctrl := gomock.NewController(t)
   123  
   124  	mockFS := fsMock.NewMockFS(ctrl)
   125  
   126  	keys := randomStrings(10)
   127  	sort.Strings(keys)
   128  	vals := randomStrings(len(keys))
   129  	reader := createSStableReader(t, keys, vals)
   130  
   131  	sut := sstable.NewPebbleSSTableRangeManagerWithNewReader(makeNewReader(reader), &NoCache{}, mockFS, crypto.SHA256)
   132  
   133  	ns := "some-ns"
   134  	sstableID := committed.ID("some-id")
   135  
   136  	iter, err := sut.NewRangeIterator(ctx, committed.Namespace(ns), sstableID)
   137  	require.NoError(t, err)
   138  	require.NotNil(t, iter)
   139  
   140  	iter.SeekGE(committed.Key(keys[len(keys)/3]))
   141  	require.NoError(t, iter.Err())
   142  
   143  	iter.Close()
   144  	require.NoError(t, iter.Err())
   145  
   146  	require.Equal(t, 1, reader.GetNumClosed())
   147  }
   148  
   149  func TestGetWriterRangeID(t *testing.T) {
   150  	ctx := context.Background()
   151  	ctrl := gomock.NewController(t)
   152  
   153  	mockFS := fsMock.NewMockFS(ctrl)
   154  
   155  	sut := sstable.NewPebbleSSTableRangeManagerWithNewReader(nil, &NoCache{}, mockFS, crypto.SHA256)
   156  
   157  	for times := 0; times < 2; times++ {
   158  		const ns = "some-ns"
   159  		mockFile := fsMock.NewMockStoredFile(ctrl)
   160  		mockFile.EXPECT().Write(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
   161  			return len(b), nil
   162  		}).AnyTimes()
   163  		mockFile.EXPECT().Sync().Return(nil).AnyTimes()
   164  		mockFile.EXPECT().Close().Return(nil).Times(1)
   165  		mockFile.EXPECT().Store(gomock.Any(), gomock.Any()).Return(nil).Times(1)
   166  		mockFS.EXPECT().Create(ctx, ns).Return(mockFile, nil).Times(1)
   167  
   168  		writer, err := sut.GetWriter(ctx, ns, nil)
   169  		require.NoError(t, err)
   170  		require.NotNil(t, writer)
   171  		err = writer.WriteRecord(committed.Record{
   172  			Key:   []byte("key"),
   173  			Value: []byte("value"),
   174  		})
   175  		require.NoError(t, err)
   176  		result, err := writer.Close()
   177  		require.NoError(t, err)
   178  		require.NotNil(t, result)
   179  		expectedID := committed.ID("1f60902cb44890618d61597673a62e44cf526f5991d9db687141218985fe60b8")
   180  		require.Equal(t, expectedID, result.RangeID, "Range ID should be kept the same based on the content")
   181  	}
   182  }