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

     1  package committed_test
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"testing"
     7  
     8  	"github.com/golang/mock/gomock"
     9  	"github.com/stretchr/testify/assert"
    10  	"github.com/stretchr/testify/require"
    11  	"github.com/treeverse/lakefs/pkg/graveler"
    12  	"github.com/treeverse/lakefs/pkg/graveler/committed"
    13  	"github.com/treeverse/lakefs/pkg/graveler/committed/mock"
    14  	"github.com/treeverse/lakefs/pkg/graveler/testutil"
    15  )
    16  
    17  func mustMarshalRange(rng committed.Range) []byte {
    18  	ret, err := committed.MarshalRange(rng)
    19  	if err != nil {
    20  		panic(err)
    21  	}
    22  	return ret
    23  }
    24  
    25  func makeValueBytesForRangeKey(key graveler.Key, idx int) []byte {
    26  	return committed.MustMarshalValue(&graveler.Value{
    27  		Identity: []byte(fmt.Sprintf("%s:%d", key, idx)),
    28  		Data:     mustMarshalRange(committed.Range{ID: ""}),
    29  	})
    30  }
    31  
    32  func makeRangeIterator(rangeIDs []graveler.Key) committed.ValueIterator {
    33  	records := make([]committed.Record, len(rangeIDs))
    34  	for i, id := range rangeIDs {
    35  		records[i] = committed.Record{
    36  			Key:   committed.Key(id),
    37  			Value: makeValueBytesForRangeKey(id, i),
    38  		}
    39  	}
    40  	return testutil.NewCommittedValueIteratorFake(records)
    41  }
    42  
    43  func makeKeys(keys ...string) []graveler.Key {
    44  	if len(keys) == 0 {
    45  		// testify doesn't think nil slices are equal to empty slices
    46  		return nil
    47  	}
    48  	ret := make([]graveler.Key, 0, len(keys))
    49  	for _, k := range keys {
    50  		ret = append(ret, graveler.Key(k))
    51  	}
    52  	return ret
    53  }
    54  
    55  type rangeKeys struct {
    56  	Name committed.ID
    57  	Keys []graveler.Key
    58  }
    59  
    60  func makeRange(r rangeKeys) committed.Range {
    61  	var ret committed.Range
    62  	if len(r.Keys) > 0 {
    63  		ret.MinKey = committed.Key(r.Keys[0])
    64  		ret.MaxKey = committed.Key(r.Keys[len(r.Keys)-1])
    65  	}
    66  	ret.ID = committed.ID(ret.MaxKey)
    67  	return ret
    68  }
    69  
    70  func makeRangeRecords(rk []rangeKeys) []committed.Record {
    71  	ret := make([]committed.Record, len(rk))
    72  	var lastKey committed.Key
    73  	for i := range rk {
    74  		rng := makeRange(rk[i])
    75  		rangeVal, err := committed.MarshalRange(rng)
    76  		if err != nil {
    77  			panic(err)
    78  		}
    79  		if rangeVal == nil {
    80  			panic("nil range")
    81  		}
    82  		key := rng.MaxKey
    83  		if len(key) == 0 {
    84  			// Empty range, MaxKey unchanged
    85  			key = lastKey
    86  		}
    87  		v := graveler.Value{Identity: key, Data: rangeVal}
    88  		vBytes := committed.MustMarshalValue(&v)
    89  		ret[i] = committed.Record{
    90  			Key:   key,
    91  			Value: vBytes,
    92  		}
    93  		lastKey = key
    94  	}
    95  	return ret
    96  }
    97  
    98  func keysByRanges(t testing.TB, it committed.Iterator) []rangeKeys {
    99  	t.Helper()
   100  	ret := make([]rangeKeys, 0)
   101  	for it.Next() {
   102  		v, p := it.Value()
   103  		require.True(t, p != nil, "iterated past end, it = %+v, %s", it, it.Err())
   104  		if v == nil {
   105  			t.Logf("new range %+v", p)
   106  			ret = append(ret, rangeKeys{Name: p.ID})
   107  		} else {
   108  			t.Logf("    element %+v", v)
   109  			p := &ret[len(ret)-1]
   110  			p.Keys = append(p.Keys, v.Key)
   111  		}
   112  	}
   113  	if err := it.Err(); err != nil {
   114  		t.Error(err)
   115  	}
   116  	return ret
   117  }
   118  
   119  func TestIterator(t *testing.T) {
   120  	namespace := committed.Namespace("ns")
   121  	tests := []struct {
   122  		Name string
   123  		PK   []rangeKeys
   124  	}{
   125  		{
   126  			Name: "empty",
   127  			PK:   []rangeKeys{},
   128  		}, {
   129  			Name: "one empty",
   130  			PK:   []rangeKeys{{Name: "", Keys: makeKeys()}},
   131  		}, {
   132  			Name: "many empty",
   133  			PK: []rangeKeys{
   134  				{Name: "", Keys: makeKeys()},
   135  				{Name: "", Keys: makeKeys()},
   136  				{Name: "", Keys: makeKeys()},
   137  			},
   138  		}, {
   139  			Name: "one",
   140  			PK:   []rangeKeys{{Name: "a3", Keys: makeKeys("a1", "a2", "a3")}},
   141  		}, {
   142  			Name: "five ranges two empty",
   143  			PK: []rangeKeys{
   144  				{Name: "a3", Keys: makeKeys("a1", "a2", "a3")},
   145  				{Name: "a3", Keys: makeKeys()},
   146  				{Name: "a3", Keys: makeKeys()},
   147  				{Name: "d1", Keys: makeKeys("d1")},
   148  				{Name: "e2", Keys: makeKeys("e1", "e2")},
   149  			},
   150  		}, {
   151  			Name: "last two empty",
   152  			PK: []rangeKeys{
   153  				{Name: "a3", Keys: makeKeys("a1", "a2", "a3")},
   154  				{Name: "a3", Keys: makeKeys()},
   155  				{Name: "a3", Keys: makeKeys()},
   156  			},
   157  		},
   158  	}
   159  
   160  	for _, tt := range tests {
   161  		t.Run(fmt.Sprintf("Iteration<%s>", tt.Name), func(t *testing.T) {
   162  			ctx := context.Background()
   163  			ctrl := gomock.NewController(t)
   164  			defer ctrl.Finish()
   165  			manager := mock.NewMockRangeManager(ctrl)
   166  
   167  			var lastKey committed.Key
   168  			for _, p := range tt.PK {
   169  				key := lastKey
   170  				if len(p.Keys) > 0 {
   171  					key = committed.Key(p.Keys[len(p.Keys)-1])
   172  				}
   173  				manager.EXPECT().
   174  					NewRangeIterator(gomock.Any(), gomock.Eq(namespace), committed.ID(key)).
   175  					Return(makeRangeIterator(p.Keys), nil)
   176  				lastKey = key
   177  			}
   178  			rangesIt := testutil.NewCommittedValueIteratorFake(makeRangeRecords(tt.PK))
   179  			pvi := committed.NewIterator(ctx, manager, namespace, rangesIt)
   180  			defer pvi.Close()
   181  			assert.Equal(t, tt.PK, keysByRanges(t, pvi))
   182  			assert.False(t, pvi.NextRange())
   183  			assert.False(t, pvi.Next())
   184  		})
   185  		t.Run(fmt.Sprintf("SeekGE<%s>", tt.Name), func(t *testing.T) {
   186  			ctx := context.Background()
   187  			ctrl := gomock.NewController(t)
   188  			defer ctrl.Finish()
   189  			manager := mock.NewMockRangeManager(ctrl)
   190  
   191  			var lastKey committed.Key
   192  			for _, p := range tt.PK {
   193  				key := lastKey
   194  				if len(p.Keys) > 0 {
   195  					key = committed.Key(p.Keys[len(p.Keys)-1])
   196  				}
   197  				manager.EXPECT().
   198  					NewRangeIterator(gomock.Any(), gomock.Eq(namespace), committed.ID(key)).
   199  					Return(makeRangeIterator(p.Keys), nil).
   200  					AnyTimes()
   201  				lastKey = key
   202  			}
   203  			rangesIt := testutil.NewCommittedValueIteratorFake(makeRangeRecords(tt.PK))
   204  			pvi := committed.NewIterator(ctx, manager, namespace, rangesIt)
   205  			defer pvi.Close()
   206  
   207  			if len(tt.PK) == 0 {
   208  				pvi.SeekGE(graveler.Key("infinity"))
   209  				if pvi.Next() {
   210  					v, r := pvi.Value()
   211  					t.Errorf("successful seeked on a nil range, value %+v,%+v", v, r)
   212  				}
   213  				return
   214  			}
   215  
   216  			for _, p := range tt.PK {
   217  				if len(p.Keys) == 0 {
   218  					continue
   219  				}
   220  				for _, k := range p.Keys {
   221  					t.Run(fmt.Sprintf("Exact:%s", string(k)), func(t *testing.T) {
   222  						pvi.SeekGE(k)
   223  						if pvi.Err() != nil {
   224  							t.Fatalf("failed to seek to %s: %s", string(k), pvi.Err())
   225  						}
   226  						if !pvi.Next() {
   227  							t.Fatalf("failed to get value at %s: %s", string(k), pvi.Err())
   228  						}
   229  						v, r := pvi.Value()
   230  						if v == nil && r != nil {
   231  							// got range - validate range
   232  							if string(k) != string(r.MinKey) {
   233  								t.Errorf("got range with MinKey %s != expected %s", string(v.Key), string(k))
   234  							}
   235  							// call next to enter range and receive value
   236  							if !pvi.Next() {
   237  								t.Fatalf("failed to get value at %s: %s", string(k), pvi.Err())
   238  							}
   239  							v, r = pvi.Value()
   240  						}
   241  						if v == nil || r == nil {
   242  							t.Fatalf("missing value-and-range %+v, %+v", v, r)
   243  						}
   244  						if string(k) != string(v.Key) {
   245  							t.Errorf("got value with ID %s != expected %s", string(v.Key), string(k))
   246  						}
   247  						if string(p.Name) != string(r.ID) {
   248  							t.Errorf("got range with ID %s != expected %s", string(r.ID), string(p.Name))
   249  						}
   250  					})
   251  				}
   252  			}
   253  		})
   254  	}
   255  }