github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/kv/bulk/kv_buf_test.go (about) 1 // Copyright 2019 The Cockroach Authors. 2 // 3 // Use of this software is governed by the Business Source License 4 // included in the file licenses/BSL.txt. 5 // 6 // As of the Change Date specified in that file, in accordance with 7 // the Business Source License, use of this software will be governed 8 // by the Apache License, Version 2.0, included in the file 9 // licenses/APL.txt. 10 11 package bulk 12 13 import ( 14 "bytes" 15 "sort" 16 "testing" 17 18 "github.com/cockroachdb/cockroach/pkg/roachpb" 19 "github.com/cockroachdb/cockroach/pkg/util/leaktest" 20 "github.com/cockroachdb/cockroach/pkg/util/randutil" 21 ) 22 23 // kvPair is a bytes -> bytes kv pair. 24 type kvPair struct { 25 key roachpb.Key 26 value []byte 27 } 28 29 func makeTestData(num int) (kvs []kvPair, totalSize int) { 30 kvs = make([]kvPair, num) 31 r, _ := randutil.NewPseudoRand() 32 alloc := make([]byte, num*500) 33 randutil.ReadTestdataBytes(r, alloc) 34 for i := range kvs { 35 if len(alloc) < 1500 { 36 const refill = 15000 37 alloc = make([]byte, refill) 38 randutil.ReadTestdataBytes(r, alloc) 39 } 40 kvs[i].key = alloc[:randutil.RandIntInRange(r, 2, 100)] 41 alloc = alloc[len(kvs[i].key):] 42 kvs[i].value = alloc[:randutil.RandIntInRange(r, 0, 1000)] 43 alloc = alloc[len(kvs[i].value):] 44 totalSize += len(kvs[i].key) + len(kvs[i].value) 45 } 46 return kvs, totalSize 47 } 48 49 func TestKvBuf(t *testing.T) { 50 defer leaktest.AfterTest(t)() 51 52 src, totalSize := makeTestData(50000) 53 54 // Write everything to our buf. 55 b := kvBuf{} 56 for i := range src { 57 if err := b.append(src[i].key, src[i].value); err != nil { 58 t.Fatal(err) 59 } 60 } 61 62 // Sanity check our buf has right size. 63 if expected, actual := len(src), b.Len(); expected != actual { 64 t.Fatalf("expected len %d got %d", expected, actual) 65 } 66 if expected, actual := totalSize+len(src)*16, b.MemSize; expected != actual { 67 t.Fatalf("expected len %d got %d", expected, actual) 68 } 69 70 // Read back what we wrote. 71 for i := range src { 72 if expected, actual := src[i].key, b.Key(i); !bytes.Equal(expected, actual) { 73 t.Fatalf("expected %s\ngot %s", expected, actual) 74 } 75 if expected, actual := src[i].value, b.Value(i); !bytes.Equal(expected, actual) { 76 t.Fatalf("expected %s\ngot %s", expected, actual) 77 } 78 } 79 // Sort both and then ensure they match. 80 sort.Slice(src, func(i, j int) bool { return bytes.Compare(src[i].key, src[j].key) < 0 }) 81 sort.Sort(&b) 82 for i := range src { 83 if expected, actual := src[i].key, b.Key(i); !bytes.Equal(expected, actual) { 84 t.Fatalf("expected %s\ngot %s", expected, actual) 85 } 86 if expected, actual := src[i].value, b.Value(i); !bytes.Equal(expected, actual) { 87 t.Fatalf("expected %s\ngot %s", expected, actual) 88 } 89 } 90 }