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 }