github.com/grailbio/base@v0.0.11/mapio/block_test.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package mapio
     6  
     7  import (
     8  	"bytes"
     9  	"math/rand"
    10  	"sort"
    11  	"testing"
    12  
    13  	fuzz "github.com/google/gofuzz"
    14  )
    15  
    16  type entry struct{ Key, Value []byte }
    17  
    18  func (e entry) Equal(f entry) bool {
    19  	return bytes.Compare(e.Key, f.Key) == 0 &&
    20  		bytes.Compare(e.Value, f.Value) == 0
    21  }
    22  
    23  func makeEntries(n int) []entry {
    24  	fz := fuzz.New()
    25  	//	fz.NumElements(1, 100)
    26  	entries := make([]entry, n)
    27  	// Fuzz manually so we can control the number of entries.
    28  	for i := range entries {
    29  		fz.Fuzz(&entries[i].Key)
    30  		fz.Fuzz(&entries[i].Value)
    31  	}
    32  	sortEntries(entries)
    33  	return entries
    34  }
    35  
    36  func sortEntries(entries []entry) {
    37  	sort.Slice(entries, func(i, j int) bool {
    38  		x := bytes.Compare(entries[i].Key, entries[j].Key)
    39  		return x < 0 || x == 0 && bytes.Compare(entries[i].Value, entries[j].Value) < 0
    40  	})
    41  }
    42  
    43  type seeker interface {
    44  	Seek(key []byte) Scanner
    45  }
    46  
    47  func testSeeker(t *testing.T, entries []entry, seeker seeker) {
    48  	t.Helper()
    49  
    50  	s := seeker.Seek(nil)
    51  	var scanned []entry
    52  	for s.Scan() {
    53  		var (
    54  			key   = make([]byte, len(s.Key()))
    55  			value = make([]byte, len(s.Value()))
    56  		)
    57  		copy(key, s.Key())
    58  		copy(value, s.Value())
    59  		scanned = append(scanned, entry{key, value})
    60  	}
    61  	if got, want := len(scanned), len(entries); got != want {
    62  		t.Fatalf("got %v, want %v", got, want)
    63  	}
    64  	isSorted := sort.SliceIsSorted(scanned, func(i, j int) bool {
    65  		return bytes.Compare(scanned[i].Key, scanned[j].Key) < 0
    66  	})
    67  	if !isSorted {
    68  		t.Error("scan returned non-sorted entries")
    69  	}
    70  	sortEntries(scanned)
    71  	sortEntries(entries)
    72  	for i := range scanned {
    73  		if !entries[i].Equal(scanned[i]) {
    74  			t.Errorf("scan: entry %d does not match", i)
    75  		}
    76  	}
    77  
    78  	// Look up keys but in a random order.
    79  	for n, i := range rand.Perm(len(entries)) {
    80  		s := seeker.Seek(entries[i].Key)
    81  		for s.Scan() {
    82  			if bytes.Compare(entries[i].Key, s.Key()) != 0 {
    83  				t.Errorf("%d: did not find key for %d", n, i)
    84  			}
    85  
    86  			// Since we may have multiple keys with the same value,
    87  			// we have to scan until we see our expected value.
    88  			if bytes.Compare(entries[i].Key, s.Key()) != 0 ||
    89  				bytes.Compare(entries[i].Value, s.Value()) == 0 {
    90  				break
    91  			}
    92  		}
    93  		if !entries[i].Equal(entry{s.Key(), s.Value()}) {
    94  			t.Errorf("%d: seek: entry %d does not match %d", n, i, bytes.Compare(entries[i].Key, s.Key()))
    95  		}
    96  	}
    97  
    98  	lastKey := entries[len(entries)-1].Key
    99  	bigKey := make([]byte, len(lastKey)+1)
   100  	copy(bigKey, lastKey)
   101  	s = seeker.Seek(bigKey)
   102  	if s.Scan() {
   103  		t.Error("scanned bigger key")
   104  	}
   105  }
   106  
   107  type blockSeeker struct{ *block }
   108  
   109  func (b *blockSeeker) Seek(key []byte) Scanner {
   110  	b.block.Seek(key)
   111  	return b
   112  }
   113  
   114  func (b *blockSeeker) Err() error { return nil }
   115  
   116  func TestBlock(t *testing.T) {
   117  	const N = 10000
   118  	entries := makeEntries(N)
   119  
   120  	buf := blockBuffer{restartInterval: 100}
   121  	for i := range entries {
   122  		buf.Append(entries[i].Key, entries[i].Value)
   123  	}
   124  	buf.Finish()
   125  
   126  	block, err := readBlock(buf.Bytes())
   127  	if err != nil {
   128  		t.Fatal(err)
   129  	}
   130  
   131  	if got, want := block.nrestart, 100; got != want {
   132  		t.Errorf("got %v, want %v", got, want)
   133  	}
   134  
   135  	testSeeker(t, entries, &blockSeeker{block})
   136  }