github.com/grailbio/base@v0.0.11/file/fsnode/iterators_test.go (about)

     1  package fsnode
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"io"
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  	"github.com/stretchr/testify/require"
    11  )
    12  
    13  func TestSliceAll(t *testing.T) {
    14  	ctx := context.Background()
    15  	nodes := []mockLeaf{{id: 0}, {id: 1}, {id: 2}}
    16  	got, err := IterateAll(ctx, NewIterator(nodes[0], nodes[1], nodes[2]))
    17  	require.NoError(t, err)
    18  	assert.Equal(t, []T{nodes[0], nodes[1], nodes[2]}, got)
    19  }
    20  
    21  func TestSliceFull(t *testing.T) {
    22  	ctx := context.Background()
    23  	nodes := []mockLeaf{{id: 0}, {id: 1}, {id: 2}}
    24  	dst := make([]T, 2)
    25  
    26  	iter := NewIterator(nodes[0], nodes[1], nodes[2])
    27  	got, err := IterateFull(ctx, iter, dst)
    28  	require.NoError(t, err)
    29  	assert.Equal(t, 2, got)
    30  	assert.Equal(t, []T{nodes[0], nodes[1]}, dst)
    31  
    32  	got, err = IterateFull(ctx, iter, dst)
    33  	assert.ErrorIs(t, err, io.EOF)
    34  	assert.Equal(t, 1, got)
    35  	assert.Equal(t, []T{nodes[2]}, dst[:1])
    36  }
    37  
    38  func TestLazy(t *testing.T) {
    39  	ctx := context.Background()
    40  	nodes := []mockLeaf{{id: 0}, {id: 1}, {id: 2}}
    41  	makeNodes := func(context.Context) ([]T, error) {
    42  		return []T{nodes[0], nodes[1], nodes[2]}, nil
    43  	}
    44  	got, err := IterateAll(ctx, NewLazyIterator(makeNodes))
    45  	require.NoError(t, err)
    46  	assert.Equal(t, []T{nodes[0], nodes[1], nodes[2]}, got)
    47  }
    48  
    49  func TestLazyErr(t *testing.T) {
    50  	ctx := context.Background()
    51  	makeNodes := func(context.Context) ([]T, error) {
    52  		return nil, errors.New("test error")
    53  	}
    54  	_, err := IterateAll(ctx, NewLazyIterator(makeNodes))
    55  	require.Error(t, err)
    56  	require.Contains(t, err.Error(), "test error")
    57  }
    58  
    59  func TestConcatAll(t *testing.T) {
    60  	ctx := context.Background()
    61  	nodes := []mockLeaf{{id: 0}, {id: 1}, {id: 2}, {id: 3}}
    62  	iter := NewConcatIterator(
    63  		NewIterator(),
    64  		NewIterator(nodes[0]),
    65  		NewIterator(),
    66  		NewIterator(),
    67  		NewIterator(nodes[1], nodes[2], nodes[3]),
    68  		NewIterator(),
    69  	)
    70  	got, err := IterateAll(ctx, iter)
    71  	require.NoError(t, err)
    72  	assert.Equal(t, []T{nodes[0], nodes[1], nodes[2], nodes[3]}, got)
    73  }
    74  
    75  func TestConcatFull(t *testing.T) {
    76  	ctx := context.Background()
    77  	nodes := []mockLeaf{{id: 0}, {id: 1}, {id: 2}, {id: 3}}
    78  	iter := NewConcatIterator(
    79  		NewIterator(),
    80  		NewIterator(nodes[0]),
    81  		NewIterator(),
    82  		NewIterator(),
    83  		NewIterator(nodes[1], nodes[2], nodes[3]),
    84  		NewIterator(),
    85  	)
    86  
    87  	var dst []T
    88  	got, err := IterateFull(ctx, iter, dst)
    89  	require.NoError(t, err)
    90  	assert.Equal(t, 0, got)
    91  
    92  	dst = make([]T, 3)
    93  	got, err = IterateFull(ctx, iter, dst[:2])
    94  	require.NoError(t, err)
    95  	assert.Equal(t, 2, got)
    96  	assert.Equal(t, []T{nodes[0], nodes[1]}, dst[:2])
    97  
    98  	got, err = IterateFull(ctx, iter, dst[:1])
    99  	require.NoError(t, err)
   100  	assert.Equal(t, 1, got)
   101  	assert.Equal(t, []T{nodes[2]}, dst[:1])
   102  
   103  	got, err = IterateFull(ctx, iter, dst)
   104  	assert.ErrorIs(t, err, io.EOF)
   105  	assert.Equal(t, 1, got)
   106  	assert.Equal(t, []T{nodes[3]}, dst[:1])
   107  }
   108  
   109  type mockLeaf struct {
   110  	Leaf
   111  	id int
   112  }