github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/slicetest/run_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  package slicetest_test
     5  
     6  import (
     7  	"math/rand"
     8  	"reflect"
     9  	"sort"
    10  	"strings"
    11  	"testing"
    12  
    13  	"github.com/grailbio/bigslice"
    14  	"github.com/grailbio/bigslice/slicetest"
    15  )
    16  
    17  var words = []string{"few", "espy", "longer", "until", "interesting",
    18  	"thus", "bason", "passage", "classes", "straighten",
    19  	"ill", "property", "combine", "promise", "Chicago",
    20  	"generally", "yellow", "per", "verb", "products",
    21  	"file", "park", "doughty", "changed", "inaureoled",
    22  	"flummoxed", "knife", "ghyll", "none", "bulwark",
    23  	"provide", "background", "purposes", "bivouac", "removed",
    24  	"jr", "adopt", "oil", "clean", "dingle",
    25  	"experience", "population", "coomb", "slightly", "encourage",
    26  	"kill", "mark", "present", "system", "standards",
    27  	"rapid", "mean", "conditions", "control", "within",
    28  	"circlet", "proper", "nothing", "craven", "case",
    29  	"day", "serious", "might", "sound", "hadn't",
    30  	"student", "rhode", "yammer", "caught", "deem",
    31  	"homes", "marry", "stands", "elect", "modern",
    32  	"activity", "servant", "too", "sport", "block",
    33  	"low", "addition", "London", "accept", "unusual",
    34  	"commercial", "grind", "books", "countries", "collect",
    35  	"these", "marriage", "narrow", "honor", "gave",
    36  	"lip", "country", "spring", "watching", "sea",
    37  }
    38  
    39  func randString(r *rand.Rand, n int) string {
    40  	var b strings.Builder
    41  	for i := 0; i < n; i++ {
    42  		b.WriteString(words[r.Intn(len(words))])
    43  	}
    44  	return b.String()
    45  }
    46  
    47  func randStringSlice(r *rand.Rand, n int) []string {
    48  	strs := make([]string, n)
    49  	for i := range strs {
    50  		strs[i] = randString(r, r.Intn(5))
    51  	}
    52  	return strs
    53  }
    54  
    55  func randIntSlice(r *rand.Rand, n int) []int {
    56  	ints := make([]int, n)
    57  	for i := range ints {
    58  		ints[i] = r.Int()
    59  	}
    60  	return ints
    61  }
    62  
    63  func TestRunAndScan(t *testing.T) {
    64  	const N = 10000
    65  	var (
    66  		r    = rand.New(rand.NewSource(0))
    67  		strs = randStringSlice(r, N)
    68  		ints = randIntSlice(r, N)
    69  	)
    70  	slice := bigslice.Const(10, strs, ints)
    71  	var (
    72  		scannedStrs []string
    73  		scannedInts []int
    74  	)
    75  	slicetest.RunAndScan(t, slice, &scannedStrs, &scannedInts)
    76  	if got, want := len(scannedStrs), N; got != want {
    77  		t.Fatalf("got %v, want %v", got, want)
    78  	}
    79  	if got, want := len(scannedInts), N; got != want {
    80  		t.Fatalf("got %v, want %v", got, want)
    81  	}
    82  	type row struct {
    83  		S string
    84  		I int
    85  	}
    86  	sortRows := func(rows []row) {
    87  		sort.Slice(rows, func(i, j int) bool {
    88  			if rows[i].S == rows[j].S {
    89  				return rows[i].I < rows[j].I
    90  			}
    91  			return rows[i].S < rows[j].S
    92  		})
    93  	}
    94  	want := make([]row, N)
    95  	for i := range want {
    96  		want[i].S = strs[i]
    97  		want[i].I = ints[i]
    98  	}
    99  	got := make([]row, N)
   100  	for i := range got {
   101  		got[i].S = scannedStrs[i]
   102  		got[i].I = scannedInts[i]
   103  	}
   104  	sortRows(got)
   105  	sortRows(want)
   106  	if !reflect.DeepEqual(got, want) {
   107  		t.Error("rows do not match")
   108  	}
   109  }