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

     1  package committed_test
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/golang/mock/gomock"
     8  	"github.com/stretchr/testify/require"
     9  	"github.com/treeverse/lakefs/pkg/catalog/testutils"
    10  	"github.com/treeverse/lakefs/pkg/graveler"
    11  	"github.com/treeverse/lakefs/pkg/graveler/committed"
    12  	"github.com/treeverse/lakefs/pkg/graveler/committed/mock"
    13  )
    14  
    15  func TestManager_WriteRange(t *testing.T) {
    16  	const (
    17  		ns         = "some-ns"
    18  		maxRecords = 4
    19  	)
    20  
    21  	writeResult := &committed.WriteResult{
    22  		RangeID:                 "someid",
    23  		First:                   committed.Key("1"),
    24  		Last:                    committed.Key("10"),
    25  		Count:                   2,
    26  		EstimatedRangeSizeBytes: 23,
    27  	}
    28  
    29  	tests := []struct {
    30  		name    string
    31  		records []*graveler.ValueRecord
    32  	}{
    33  		{
    34  			name: "iterator_exhausted",
    35  			records: []*graveler.ValueRecord{
    36  				{Key: graveler.Key("1"), Value: &graveler.Value{}},
    37  				{Key: graveler.Key("2"), Value: &graveler.Value{}},
    38  			},
    39  		},
    40  		{
    41  			name: "break_at_key",
    42  			records: []*graveler.ValueRecord{
    43  				{Key: graveler.Key("1"), Value: &graveler.Value{}},
    44  				{Key: graveler.Key("2"), Value: &graveler.Value{}},
    45  				{Key: graveler.Key("3"), Value: &graveler.Value{}},
    46  				{Key: graveler.Key("4"), Value: &graveler.Value{}},
    47  				{Key: graveler.Key("5"), Value: &graveler.Value{}},
    48  			},
    49  		},
    50  	}
    51  	for _, tt := range tests {
    52  		t.Run(tt.name, func(t *testing.T) {
    53  			ctrl := gomock.NewController(t)
    54  			metarangeManager := mock.NewMockMetaRangeManager(ctrl)
    55  			rangeManager := mock.NewMockRangeManager(ctrl)
    56  			rangeWriter := mock.NewMockRangeWriter(ctrl)
    57  
    58  			rangeWriter.EXPECT().Abort().Return(nil)
    59  			rangeManager.EXPECT().GetWriter(context.Background(), committed.Namespace(ns), nil).Return(rangeWriter, nil)
    60  
    61  			sut := committed.NewCommittedManager(metarangeManager, rangeManager, params)
    62  
    63  			times := 0
    64  			expectedTimes := _min(len(tt.records), maxRecords)
    65  			rangeWriter.EXPECT().WriteRecord(gomock.Any()).Return(nil).Times(expectedTimes)
    66  			rangeWriter.EXPECT().ShouldBreakAtKey(gomock.Any(), gomock.Any()).
    67  				DoAndReturn(func(interface{}, interface{}) bool { times++; return times == maxRecords }).Times(expectedTimes)
    68  			rangeWriter.EXPECT().Close().Return(writeResult, nil)
    69  			rangeWriter.EXPECT().SetMetadata(committed.MetadataTypeKey, committed.MetadataRangesType)
    70  
    71  			it := testutils.NewFakeValueIterator(tt.records)
    72  			rangeInfo, err := sut.WriteRange(context.Background(), ns, it)
    73  			require.NoError(t, err)
    74  			require.Equal(t, &graveler.RangeInfo{
    75  				ID:                      graveler.RangeID(writeResult.RangeID),
    76  				MinKey:                  graveler.Key(writeResult.First),
    77  				MaxKey:                  graveler.Key(writeResult.Last),
    78  				Count:                   writeResult.Count,
    79  				EstimatedRangeSizeBytes: writeResult.EstimatedRangeSizeBytes,
    80  			}, rangeInfo)
    81  		})
    82  	}
    83  }
    84  
    85  // _min - helper function to return the minimum of two integers
    86  // TODO(barak): replace with builtin min function after upgrade to go 1.21
    87  func _min(a int, b int) int {
    88  	if a < b {
    89  		return a
    90  	}
    91  	return b
    92  }
    93  
    94  func TestManager_WriteMetaRange(t *testing.T) {
    95  	const (
    96  		ns = "some-ns"
    97  	)
    98  
    99  	expectedMetarangeID := graveler.MetaRangeID("some-id")
   100  
   101  	tests := []struct {
   102  		name    string
   103  		records []*graveler.RangeInfo
   104  	}{
   105  		//{
   106  		//	name: "in_order",
   107  		//	records: []*graveler.RangeInfo{
   108  		//		{ID: "id1", MinKey: graveler.Key("1"), MaxKey: graveler.Key("11")},
   109  		//		{ID: "id2", MinKey: graveler.Key("2"), MaxKey: graveler.Key("22")},
   110  		//	},
   111  		//},
   112  		{
   113  			name: "wrong_order",
   114  			records: []*graveler.RangeInfo{
   115  				{ID: "id1", MinKey: graveler.Key("1"), MaxKey: graveler.Key("11")},
   116  				{ID: "id3", MinKey: graveler.Key("3"), MaxKey: graveler.Key("33")},
   117  				{ID: "id2", MinKey: graveler.Key("2"), MaxKey: graveler.Key("22")},
   118  			},
   119  		},
   120  	}
   121  	for _, tt := range tests {
   122  		t.Run(tt.name, func(t *testing.T) {
   123  			ctrl := gomock.NewController(t)
   124  			metarangeManager := mock.NewMockMetaRangeManager(ctrl)
   125  			rangeManager := mock.NewMockRangeManager(ctrl)
   126  			metarangeWriter := mock.NewMockMetaRangeWriter(ctrl)
   127  
   128  			minKey := ""
   129  			metarangeManager.EXPECT().NewWriter(context.Background(), graveler.StorageNamespace(ns), nil).Return(metarangeWriter)
   130  			metarangeWriter.EXPECT().WriteRange(gomock.Any()).Return(nil).
   131  				DoAndReturn(func(info committed.Range) error {
   132  					if string(info.MinKey) < minKey {
   133  						t.Fatalf("record should be sorted ascending - previous minKey '%s', current '%s'", minKey, info.MinKey)
   134  					}
   135  					minKey = string(info.MinKey)
   136  					return nil
   137  				}).Times(len(tt.records))
   138  			metarangeWriter.EXPECT().Close(gomock.Any()).Return(&expectedMetarangeID, nil)
   139  			metarangeWriter.EXPECT().Abort().Return(nil)
   140  			sut := committed.NewCommittedManager(metarangeManager, rangeManager, params)
   141  
   142  			actualMetarangeID, err := sut.WriteMetaRange(context.Background(), ns, tt.records)
   143  			require.NoError(t, err)
   144  			require.Equal(t, &graveler.MetaRangeInfo{ID: expectedMetarangeID}, actualMetarangeID)
   145  		})
   146  	}
   147  }