github.com/grafana/pyroscope@v1.18.0/pkg/iter/batch_test.go (about)

     1  package iter
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  
     8  	"github.com/samber/lo"
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func TestReadBatch(t *testing.T) {
    13  	ctx := context.Background()
    14  
    15  	require.Error(t, ReadBatch(ctx, NewSliceIterator(lo.Times(20, func(i int) int { return i })), 10,
    16  		func(context.Context, []int) error {
    17  			return errors.New("foo")
    18  		}))
    19  
    20  	var batches [][]int
    21  	require.NoError(t, ReadBatch(ctx, NewSliceIterator(lo.Times(20, func(i int) int { return i })), 10,
    22  		func(_ context.Context, batch []int) error {
    23  			c := make([]int, len(batch))
    24  			copy(c, batch)
    25  			batches = append(batches, c)
    26  			return nil
    27  		}))
    28  	require.Equal(t, [][]int{lo.Times(10, func(i int) int { return i }), lo.Times(10, func(i int) int { return 10 + i })}, batches)
    29  
    30  	batches = nil
    31  	require.NoError(t, ReadBatch(ctx, NewSliceIterator(lo.Times(20, func(i int) int { return i })), 11,
    32  		func(_ context.Context, batch []int) error {
    33  			c := make([]int, len(batch))
    34  			copy(c, batch)
    35  			batches = append(batches, c)
    36  			return nil
    37  		}))
    38  	require.Equal(t, [][]int{lo.Times(11, func(i int) int { return i }), lo.Times(9, func(i int) int { return 11 + i })}, batches)
    39  }