github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/example/max_test.go

// Copyright 2020 GRAIL, Inc. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

package example

import (
	"context"
	"math/rand"
	"reflect"
	"testing"
	"testing/quick"

	"github.com/grailbio/bigslice"
	"github.com/grailbio/bigslice/slicetest"
)

// TestIntMax verifies that IntMax works on a trivial input. Its main purpose
// is to illustrate the basic test flow and usage of slicetest.
func TestIntMax(t *testing.T) {
	slice := bigslice.Const(2,
		[]int{0, 0, 0, 1, 1, 2},
		[]int{6, 5, 4, 3, 2, 1},
	)
	slice = IntMax(slice)
	scanner := slicetest.Run(t, slice)
	var (
		key int
		val int
		got = make(map[int]int)
	)
	for scanner.Scan(context.Background(), &key, &val) {
		got[key] = val
	}
	if err := scanner.Err(); err != nil {
		t.Fatal(err)
	}
	want := map[int]int{
		0: 6,
		1: 3,
		2: 1,
	}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("got %v, want %v", got, want)
	}
}

// intMaxTestCase is a test case for our property test. Random cases will be
// generated by testing.quick.
type intMaxTestCase struct {
	numShards int
	valsByKey map[int][]int
}

// Generate implements the quick.Generator interface.
func (intMaxTestCase) Generate(r *rand.Rand, size int) reflect.Value {
	valsByKey := make(map[int][]int)
	for i := 0; i < size; i++ {
		for j := 0; j < r.Intn(size); j++ {
			valsByKey[i] = append(valsByKey[i], r.Int())
		}
	}
	return reflect.ValueOf(intMaxTestCase{
		numShards: r.Intn(size) + 1, // At least 1 shard.
		valsByKey: valsByKey,
	})
}

// TestIntMaxProperties is a more sophisticated property test of IntMax. It
// illustrates the integration of property-based testing, which is likely useful
// for asserting invariants of large-scale processing, with slice testing
// machinery.
func TestIntMaxProperties(t *testing.T) {
	f := func(c intMaxTestCase) bool {
		var (
			ks []int
			vs []int
		)
		for k, kVals := range c.valsByKey {
			for _, v := range kVals {
				ks = append(ks, k)
				vs = append(vs, v)
			}
		}
		slice := bigslice.Const(c.numShards, ks, vs)
		slice = IntMax(slice)
		scanner := slicetest.Run(t, slice)
		var (
			k   int
			max int
			got = make(map[int]int)
		)
		for scanner.Scan(context.Background(), &k, &max) {
			// Each key exists in the input.
			inVs, ok := c.valsByKey[k]
			if !ok {
				t.Logf("key not in input: %v", k)
				return false
			}
			// The maximum value exists in the input for the key.
			var found bool
			for _, inV := range inVs {
				if max == inV {
					found = true
				}
			}
			if !found {
				t.Logf("value not found key inputs for key %v: %v", k, max)
				return false
			}
			// No input for the key is greater than the computed maximum.
			for _, inV := range inVs {
				if max < inV {
					return false
				}
			}
			// No key is duplicated.
			if _, ok = got[k]; ok {
				t.Logf("duplicate k: %v", k)
				return false
			}
			got[k] = max
		}

		return true
	}
	c := quick.Config{MaxCount: 10}
	if err := quick.Check(f, &c); err != nil {
		t.Error(err)
	}
}