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

     1  package sstable_test
     2  
     3  import (
     4  	"context"
     5  	"crypto/sha256"
     6  	"sort"
     7  	"testing"
     8  
     9  	"github.com/golang/mock/gomock"
    10  	"github.com/stretchr/testify/require"
    11  	"github.com/thanhpk/randstr"
    12  	"github.com/treeverse/lakefs/pkg/graveler/committed"
    13  	"github.com/treeverse/lakefs/pkg/graveler/sstable"
    14  	"github.com/treeverse/lakefs/pkg/pyramid/mock"
    15  )
    16  
    17  func TestWriter(t *testing.T) {
    18  	ctx := context.Background()
    19  	ctrl := gomock.NewController(t)
    20  	mockFS := mock.NewMockFS(ctrl)
    21  	defer ctrl.Finish()
    22  	ns := committed.Namespace("some-namespace")
    23  
    24  	// create the mock file with the matching file-system
    25  	mockFile := mock.NewMockStoredFile(ctrl)
    26  	mockFile.EXPECT().Close().Return(nil).Times(1)
    27  	mockFS.EXPECT().Create(gomock.Any(), string(ns)).Return(mockFile, nil)
    28  
    29  	writes := 500
    30  	dw, err := sstable.NewDiskWriter(ctx, mockFS, ns, sha256.New(), nil)
    31  	require.NoError(t, err)
    32  	require.NotNil(t, dw)
    33  
    34  	keys := randomStrings(writes)
    35  	sort.Strings(keys)
    36  	var f string
    37  
    38  	// expect the specific write file actions
    39  	mockFile.EXPECT().Write(gomock.Any()).DoAndReturn(
    40  		func(b []byte) (int, error) {
    41  			return len(b), nil
    42  		}).MinTimes(1)
    43  	mockFile.EXPECT().Sync().Return(nil).AnyTimes()
    44  	mockFile.EXPECT().Store(gomock.Any(), gomock.Any()).DoAndReturn(
    45  		func(_ context.Context, filename string) error {
    46  			f = filename
    47  			return nil
    48  		}).AnyTimes()
    49  
    50  	// Do the actual writing
    51  	for i := 0; i < writes; i++ {
    52  		err = dw.WriteRecord(committed.Record{
    53  			Key:   []byte(keys[i]),
    54  			Value: []byte("some-data"),
    55  		})
    56  		require.NoError(t, err)
    57  	}
    58  
    59  	// Close and assert the result
    60  	wr, err := dw.Close()
    61  	require.NoError(t, err)
    62  	require.NotNil(t, wr)
    63  	require.Equal(t, writes, wr.Count)
    64  	require.Equal(t, keys[0], string(wr.First))
    65  	require.Equal(t, keys[writes-1], string(wr.Last))
    66  	require.Equal(t, committed.ID(f), wr.RangeID)
    67  }
    68  
    69  func TestWriterAbort(t *testing.T) {
    70  	ctx := context.Background()
    71  	ctrl := gomock.NewController(t)
    72  	mockFS := mock.NewMockFS(ctrl)
    73  	defer ctrl.Finish()
    74  	ns := committed.Namespace("some-namespace")
    75  
    76  	// create the mock file with the matching file-system
    77  	mockFile := mock.NewMockStoredFile(ctrl)
    78  	mockFile.EXPECT().Abort(gomock.Any()).Return(nil).Times(1)
    79  	mockFile.EXPECT().Close().Return(nil).Times(1)
    80  	mockFS.EXPECT().Create(gomock.Any(), string(ns)).Return(mockFile, nil)
    81  
    82  	dw, err := sstable.NewDiskWriter(ctx, mockFS, ns, sha256.New(), nil)
    83  	require.NoError(t, err)
    84  	require.NotNil(t, dw)
    85  
    86  	// expect the specific write file actions
    87  	mockFile.EXPECT().Write(gomock.Any()).DoAndReturn(
    88  		func(b []byte) (int, error) {
    89  			return len(b), nil
    90  		}).Times(1)
    91  	mockFile.EXPECT().Sync().Return(nil).AnyTimes()
    92  
    93  	// Do the actual writing
    94  	err = dw.WriteRecord(committed.Record{
    95  		Key:   []byte("key-1"),
    96  		Value: []byte("some-data"),
    97  	})
    98  	require.NoError(t, err)
    99  
   100  	// Abort
   101  	require.NoError(t, dw.Abort())
   102  }
   103  
   104  func TestWriterAbortAfterClose(t *testing.T) {
   105  	ctx := context.Background()
   106  	ctrl := gomock.NewController(t)
   107  	mockFS := mock.NewMockFS(ctrl)
   108  	defer ctrl.Finish()
   109  	ns := committed.Namespace("some-namespace")
   110  
   111  	// create the mock file with the matching file-system
   112  	mockFile := mock.NewMockStoredFile(ctrl)
   113  	mockFile.EXPECT().Close().Return(nil).Times(1)
   114  	mockFS.EXPECT().Create(gomock.Any(), string(ns)).Return(mockFile, nil)
   115  	// expect any write file actions
   116  	mockFile.EXPECT().Write(gomock.Any()).DoAndReturn(func(b []byte) (int, error) { return len(b), nil }).Times(1)
   117  	mockFile.EXPECT().Sync().Return(nil).Times(1)
   118  	mockFile.EXPECT().Store(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, filename string) error { return nil }).Times(1)
   119  
   120  	// Create writer
   121  	dw, err := sstable.NewDiskWriter(ctx, mockFS, ns, sha256.New(), nil)
   122  	require.NoError(t, err)
   123  	require.NotNil(t, dw)
   124  
   125  	// Write something writing
   126  	err = dw.WriteRecord(committed.Record{
   127  		Key:   []byte("key-1"),
   128  		Value: []byte("some-data"),
   129  	})
   130  	require.NoError(t, err)
   131  
   132  	// Close
   133  	result, err := dw.Close()
   134  	require.NoError(t, err)
   135  	require.NotNil(t, result)
   136  
   137  	// Abort
   138  	err = dw.Abort()
   139  	require.NoError(t, err)
   140  }
   141  
   142  func randomStrings(writes int) []string {
   143  	var keys []string
   144  	for i := 0; i < writes; i++ {
   145  		keys = append(keys, randstr.String(20, "abcdefghijklmnopqrstuvwyz0123456789"))
   146  	}
   147  	return keys
   148  }