github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/sliceio/spiller_test.go (about)

     1  // Copyright 2019 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 sliceio
     6  
     7  import (
     8  	"context"
     9  	"testing"
    10  
    11  	fuzz "github.com/google/gofuzz"
    12  	"github.com/grailbio/bigslice/frame"
    13  )
    14  
    15  func TestSpiller(t *testing.T) {
    16  	const n = 100
    17  	var (
    18  		fz = fuzz.NewWithSeed(123)
    19  		f1 = fuzzFrame(fz, n/2, typeOfString, typeOfInt)
    20  	)
    21  	spill, err := NewSpiller("test")
    22  	if err != nil {
    23  		t.Fatal(err)
    24  	}
    25  	defer func() {
    26  		if err = spill.Cleanup(); err != nil {
    27  			t.Fatal(err)
    28  		}
    29  	}()
    30  	if _, err = spill.Spill(f1); err != nil {
    31  		t.Fatal(err)
    32  	}
    33  	if _, err = spill.Spill(f1); err != nil {
    34  		t.Fatal(err)
    35  	}
    36  
    37  	f12 := frame.Make(f1, 2*n, 2*n)
    38  	frame.Copy(f12, f1)
    39  	frame.Copy(f12.Slice(n/2, n), f1)
    40  
    41  	readers, err := spill.Readers()
    42  	if err != nil {
    43  		t.Fatal(err)
    44  	}
    45  	if got, want := len(readers), 2; got != want {
    46  		t.Fatalf("got %v, want %v", got, want)
    47  	}
    48  	reader := MultiReader(readers...)
    49  	defer reader.Close()
    50  	m, err := ReadFull(context.Background(), reader, f12.Slice(n, 2*n))
    51  	if err != nil {
    52  		t.Fatal(err)
    53  	}
    54  	if got, want := n, m; got != want {
    55  		t.Fatalf("got %v, want %v", got, want)
    56  	}
    57  	m, err = reader.Read(context.Background(), frame.Make(f12, 1, 1))
    58  	if got, want := err, EOF; got != want {
    59  		t.Errorf("got %v, want %v", got, want)
    60  	}
    61  	if got, want := m, 0; got != want {
    62  		t.Errorf("got %v, want %v", got, want)
    63  	}
    64  	for i := 0; i < n; i++ {
    65  		if f12.Less(i, n+i) || f12.Less(n+i, i) {
    66  			t.Fatalf("not equal: mismatch at %d", i)
    67  		}
    68  	}
    69  
    70  }