github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/pkg/genutil/slicez/chunking_test.go (about)

     1  package slicez
     2  
     3  import (
     4  	"fmt"
     5  	"math"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/require"
     9  )
    10  
    11  func TestForEachChunk(t *testing.T) {
    12  	t.Parallel()
    13  
    14  	for _, datasize := range []int{0, 1, 5, 10, 50, 100, 250} {
    15  		datasize := datasize
    16  		for _, chunksize := range []uint16{1, 2, 3, 5, 10, 50} {
    17  			chunksize := chunksize
    18  			t.Run(fmt.Sprintf("test-%d-%d", datasize, chunksize), func(t *testing.T) {
    19  				t.Parallel()
    20  
    21  				data := make([]int, 0, datasize)
    22  				for i := 0; i < datasize; i++ {
    23  					data = append(data, i)
    24  				}
    25  
    26  				found := make([]int, 0, datasize)
    27  				ForEachChunk(data, chunksize, func(items []int) {
    28  					found = append(found, items...)
    29  					require.True(t, len(items) <= int(chunksize))
    30  					require.True(t, len(items) > 0)
    31  				})
    32  				require.Equal(t, data, found)
    33  			})
    34  		}
    35  	}
    36  }
    37  
    38  func TestForEachChunkOverflowPanic(t *testing.T) {
    39  	t.Parallel()
    40  
    41  	datasize := math.MaxUint16
    42  	chunksize := uint16(50)
    43  	data := make([]int, 0, datasize)
    44  	for i := 0; i < datasize; i++ {
    45  		data = append(data, i)
    46  	}
    47  
    48  	found := make([]int, 0, datasize)
    49  	ForEachChunk(data, chunksize, func(items []int) {
    50  		found = append(found, items...)
    51  		require.True(t, len(items) <= int(chunksize))
    52  		require.True(t, len(items) > 0)
    53  	})
    54  
    55  	require.Equal(t, data, found)
    56  }
    57  
    58  func TestForEachChunkOverflowIncorrect(t *testing.T) {
    59  	t.Parallel()
    60  
    61  	chunksize := uint16(50)
    62  	for _, datasize := range []int{math.MaxUint16 + int(chunksize), 10_000_000} {
    63  		datasize := datasize
    64  		t.Run(fmt.Sprintf("test-%d-%d", datasize, chunksize), func(t *testing.T) {
    65  			t.Parallel()
    66  
    67  			data := make([]int, 0, datasize)
    68  			for i := 0; i < datasize; i++ {
    69  				data = append(data, i)
    70  			}
    71  
    72  			found := make([]int, 0, datasize)
    73  			ForEachChunk(data, chunksize, func(items []int) {
    74  				found = append(found, items...)
    75  				require.True(t, len(items) <= int(chunksize))
    76  				require.True(t, len(items) > 0)
    77  			})
    78  
    79  			require.Equal(t, data, found)
    80  		})
    81  	}
    82  }
    83  
    84  func TestForEachChunkUntil(t *testing.T) {
    85  	for _, datasize := range []int{0, 1, 5, 10, 50, 100, 250} {
    86  		datasize := datasize
    87  		for _, chunksize := range []uint16{1, 2, 3, 5, 10, 50} {
    88  			chunksize := chunksize
    89  			t.Run(fmt.Sprintf("test-%d-%d", datasize, chunksize), func(t *testing.T) {
    90  				var data []int
    91  				for i := 0; i < datasize; i++ {
    92  					data = append(data, i)
    93  				}
    94  
    95  				var found []int
    96  				ok, err := ForEachChunkUntil(data, chunksize, func(items []int) (bool, error) {
    97  					found = append(found, items...)
    98  					require.True(t, len(items) <= int(chunksize))
    99  					require.True(t, len(items) > 0)
   100  					return true, nil
   101  				})
   102  				require.True(t, ok)
   103  				require.NoError(t, err)
   104  				require.Equal(t, data, found)
   105  			})
   106  		}
   107  	}
   108  }
   109  
   110  func TestForEachChunkUntilCancels(t *testing.T) {
   111  	ok, err := ForEachChunkUntil([]int{1, 2, 3, 4}, 2, func(items []int) (bool, error) {
   112  		require.Equal(t, []int{1, 2}, items)
   113  		return false, nil
   114  	})
   115  	require.False(t, ok)
   116  	require.NoError(t, err)
   117  }
   118  
   119  func TestForEachChunkUntilErrors(t *testing.T) {
   120  	ok, err := ForEachChunkUntil([]int{1, 2, 3, 4}, 2, func(items []int) (bool, error) {
   121  		require.Equal(t, []int{1, 2}, items)
   122  		return true, fmt.Errorf("some error")
   123  	})
   124  	require.False(t, ok)
   125  	require.Error(t, err)
   126  }