github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/slicetest/run.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 slicetest provides utilities for testing Bigslice user code.
     6  // The utilities here are generally not optimized for performance or
     7  // robustness; they are strictly intended for unit testing.
     8  package slicetest
     9  
    10  import (
    11  	"context"
    12  	"reflect"
    13  	"testing"
    14  
    15  	"github.com/grailbio/bigslice"
    16  	"github.com/grailbio/bigslice/exec"
    17  	"github.com/grailbio/bigslice/sliceio"
    18  )
    19  
    20  // Run evaluates the provided slice in local execution mode,
    21  // returning a scanner for the result. Errors are reported as fatal
    22  // to the provided t instance. Run is intended for unit testing of
    23  // Slice implementations.
    24  func Run(t *testing.T, slice bigslice.Slice) *sliceio.Scanner {
    25  	t.Helper()
    26  	ctx := context.Background()
    27  	fn := bigslice.Func(func() bigslice.Slice { return slice })
    28  	sess := exec.Start(exec.Local)
    29  	res, err := sess.Run(ctx, fn)
    30  	if err != nil {
    31  		t.Fatal(err)
    32  	}
    33  	return res.Scanner()
    34  }
    35  
    36  // RunErr evaluates the provided slice in local execution mode
    37  // and returns the error, if any.
    38  func RunErr(slice bigslice.Slice) error {
    39  	ctx := context.Background()
    40  	fn := bigslice.Func(func() bigslice.Slice { return slice })
    41  	sess := exec.Start(exec.Local)
    42  	_, err := sess.Run(ctx, fn)
    43  	return err
    44  }
    45  
    46  // ScanAll scans all entries from the scanner into the provided
    47  // columns, which must be pointers to slices of the correct column
    48  // types. For example, to read all values for a Slice<int, string>:
    49  //
    50  //	var (
    51  //		ints []int
    52  //		strings []string
    53  //	)
    54  //	ScanAll(test, scan, &ints, &strings)
    55  //
    56  // Errors are reported as fatal to the provided t instance.
    57  func ScanAll(t *testing.T, scan *sliceio.Scanner, cols ...interface{}) {
    58  	t.Helper()
    59  	vs := make([]reflect.Value, len(cols))
    60  	elemTypes := make([]reflect.Type, len(cols))
    61  	for i := range vs {
    62  		vs[i] = reflect.Indirect(reflect.ValueOf(cols[i]))
    63  		vs[i].Set(vs[i].Slice(0, 0))
    64  		elemTypes[i] = vs[i].Type().Elem()
    65  	}
    66  	ctx := context.Background()
    67  	args := make([]interface{}, len(cols))
    68  	for n := 0; ; n++ {
    69  		for i := range vs {
    70  			vs[i].Set(reflect.Append(vs[i], reflect.Zero(elemTypes[i])))
    71  			args[i] = vs[i].Index(n).Addr().Interface()
    72  		}
    73  		if !scan.Scan(ctx, args...) {
    74  			for i := range vs {
    75  				vs[i].Set(vs[i].Slice(0, n))
    76  			}
    77  			break
    78  		}
    79  	}
    80  	if err := scan.Err(); err != nil {
    81  		t.Fatal(err)
    82  	}
    83  }
    84  
    85  // RunAndScan evaluates the provided slice and scans its results into
    86  // the provided slice pointers. Errors are reported as fatal to the provided
    87  // t instance.
    88  func RunAndScan(t *testing.T, slice bigslice.Slice, cols ...interface{}) {
    89  	t.Helper()
    90  	scanner := Run(t, slice)
    91  	defer scanner.Close()
    92  	ScanAll(t, scanner, cols...)
    93  }