github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/sliceio/reader_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 sliceio
     6  
     7  import (
     8  	"context"
     9  	"errors"
    10  	"reflect"
    11  	"testing"
    12  
    13  	fuzz "github.com/google/gofuzz"
    14  	"github.com/grailbio/bigslice/frame"
    15  )
    16  
    17  func TestFrameReader(t *testing.T) {
    18  	const N = 1000
    19  	var (
    20  		fz  = fuzz.NewWithSeed(12345)
    21  		f   = fuzzFrame(fz, N, typeOfString)
    22  		r   = FrameReader(f)
    23  		out = frame.Make(f, N, N)
    24  		ctx = context.Background()
    25  	)
    26  	n, err := ReadFull(ctx, r, out)
    27  	if err != nil && err != EOF {
    28  		t.Fatal(err)
    29  	}
    30  	if got, want := n, N; got != want {
    31  		t.Fatalf("got %v, want %v", got, want)
    32  	}
    33  	if err == nil {
    34  		n, err := ReadFull(ctx, r, frame.Make(f, 1, 1))
    35  		if got, want := err, EOF; got != want {
    36  			t.Errorf("got %v, want %v", got, want)
    37  		}
    38  		if got, want := n, 0; got != want {
    39  			t.Errorf("got %v, want %v", got, want)
    40  		}
    41  	}
    42  
    43  	if !reflect.DeepEqual(f.Interface(0).([]string), out.Interface(0).([]string)) {
    44  		t.Error("frames do not match")
    45  	}
    46  }
    47  
    48  // TestMultiReaderClose verifies that (*multiReader).Close closes all of the
    49  // comprising readers.
    50  func TestMultiReaderClose(t *testing.T) {
    51  	const NReaders = 10
    52  	var (
    53  		fz        = fuzz.NewWithSeed(12345)
    54  		readers   = make([]ReadCloser, NReaders)
    55  		numClosed int
    56  	)
    57  	for i := range readers {
    58  		f := fuzzFrame(fz, 1000, typeOfString)
    59  		closeFunc := func() error {
    60  			numClosed++
    61  			return nil
    62  		}
    63  		readers[i] = ReaderWithCloseFunc{FrameReader(f), closeFunc}
    64  	}
    65  	r := MultiReader(readers...)
    66  	r.Close()
    67  	if got, want := numClosed, NReaders; got != want {
    68  		t.Errorf("got %v, want %v", got, want)
    69  	}
    70  }
    71  
    72  // TestMultiReaderClose verifies that (*multiReader).Close closes all of the
    73  // comprising readers, even if not all readers have been exhausted.
    74  func TestMultiReaderErrClose(t *testing.T) {
    75  	const NReaders = 10
    76  	var (
    77  		fz        = fuzz.NewWithSeed(12345)
    78  		readers   = make([]ReadCloser, NReaders)
    79  		numClosed int
    80  	)
    81  	for i := range readers {
    82  		closeFunc := func() error {
    83  			numClosed++
    84  			return nil
    85  		}
    86  		// One of the readers in the middle returns an error.
    87  		if i == 3 {
    88  			readers[i] = ReaderWithCloseFunc{ErrReader(errors.New("some error")), closeFunc}
    89  			continue
    90  		}
    91  		f := fuzzFrame(fz, 1000, typeOfString)
    92  		readers[i] = ReaderWithCloseFunc{FrameReader(f), closeFunc}
    93  	}
    94  	r := MultiReader(readers...)
    95  	r.Close()
    96  	// Make sure all readers are closed, despite error in the middle.
    97  	if got, want := numClosed, NReaders; got != want {
    98  		t.Errorf("got %v, want %v", got, want)
    99  	}
   100  }
   101  
   102  // FuzzFrame creates a fuzzed frame of length n, where columns
   103  // have the provided types.
   104  func fuzzFrame(fz *fuzz.Fuzzer, n int, types ...reflect.Type) frame.Frame {
   105  	cols := make([]interface{}, len(types))
   106  	for i := range cols {
   107  		v := reflect.MakeSlice(reflect.SliceOf(types[i]), n, n)
   108  		vp := reflect.New(types[i])
   109  		for j := 0; j < n; j++ {
   110  			fz.Fuzz(vp.Interface())
   111  			v.Index(j).Set(vp.Elem())
   112  		}
   113  		cols[i] = v.Interface()
   114  	}
   115  	return frame.Slices(cols...)
   116  }