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  }