github.com/grailbio/base@v0.0.11/file/addfs/per_node_test.go (about)

     1  package addfs
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sort"
     7  	"strings"
     8  	"testing"
     9  
    10  	"github.com/grailbio/base/file/fsnode"
    11  	. "github.com/grailbio/base/file/fsnode/fsnodetesting"
    12  	"github.com/stretchr/testify/assert"
    13  	"github.com/stretchr/testify/require"
    14  )
    15  
    16  func TestPerNodeFuncs(t *testing.T) {
    17  	ctx := context.Background()
    18  	root := func() Parent {
    19  		return Parent{
    20  			"dir0": Parent{},
    21  			"dir1": Parent{
    22  				"dir10": Parent{
    23  					"a": []byte("content dir10/a"),
    24  					"b": []byte("content dir10/b"),
    25  				},
    26  				"a": []byte("content dir1/a"),
    27  				"b": []byte("content dir1/b"),
    28  			},
    29  		}
    30  	}
    31  	t.Run("basic", func(t *testing.T) {
    32  		root := root()
    33  		n := MakeT(t, "", root).(fsnode.Parent)
    34  		n = ApplyPerNodeFuncs(n,
    35  			NewPerNodeFunc(
    36  				func(ctx context.Context, node fsnode.T) ([]fsnode.T, error) {
    37  					switch n := node.(type) {
    38  					case fsnode.Parent:
    39  						iter := n.Children()
    40  						defer func() { assert.NoError(t, iter.Close(ctx)) }()
    41  						children, err := fsnode.IterateAll(ctx, iter)
    42  						assert.NoError(t, err)
    43  						var names []string
    44  						for _, child := range children {
    45  							names = append(names, child.Info().Name())
    46  						}
    47  						sort.Strings(names)
    48  						return []fsnode.T{
    49  							fsnode.ConstLeaf(fsnode.NewRegInfo("children names"), []byte(strings.Join(names, ","))),
    50  						}, nil
    51  					case fsnode.Leaf:
    52  						return []fsnode.T{
    53  							fsnode.ConstLeaf(fsnode.NewRegInfo("copy"), nil), // Will be overwritten.
    54  						}, nil
    55  					}
    56  					require.Failf(t, "invalid node type", "node: %T", node)
    57  					panic("unreachable")
    58  				},
    59  			),
    60  			NewPerNodeFunc(
    61  				func(ctx context.Context, node fsnode.T) ([]fsnode.T, error) {
    62  					switch n := node.(type) {
    63  					case fsnode.Parent:
    64  						return nil, nil
    65  					case fsnode.Leaf:
    66  						return []fsnode.T{
    67  							fsnode.ConstLeaf(fsnode.NewRegInfo("copy"), LeafReadAll(ctx, t, n)),
    68  						}, nil
    69  					}
    70  					require.Failf(t, "invalid node type", "node: %T", node)
    71  					panic("unreachable")
    72  				},
    73  			),
    74  		)
    75  		got := Walker{}.WalkContents(ctx, t, n)
    76  		want := Parent{
    77  			"...": Parent{
    78  				"dir0": Parent{"children names": []byte("")},
    79  				"dir1": Parent{"children names": []byte("a,b,dir10")},
    80  			},
    81  			"dir0": Parent{
    82  				"...": Parent{},
    83  			},
    84  			"dir1": Parent{
    85  				"...": Parent{
    86  					"dir10": Parent{"children names": []byte("a,b")},
    87  					"a":     Parent{"copy": []byte("content dir1/a")},
    88  					"b":     Parent{"copy": []byte("content dir1/b")},
    89  				},
    90  				"dir10": Parent{
    91  					"...": Parent{
    92  						"a": Parent{"copy": []byte("content dir10/a")},
    93  						"b": Parent{"copy": []byte("content dir10/b")},
    94  					},
    95  					"a": []byte("content dir10/a"),
    96  					"b": []byte("content dir10/b"),
    97  				},
    98  				"a": []byte("content dir1/a"),
    99  				"b": []byte("content dir1/b"),
   100  			},
   101  		}
   102  		assert.Equal(t, want, got)
   103  	})
   104  	t.Run("lazy", func(t *testing.T) {
   105  		root := root()
   106  		n := MakeT(t, "", root).(fsnode.Parent)
   107  		n = ApplyPerNodeFuncs(n, NewPerNodeFunc(
   108  			func(_ context.Context, node fsnode.T) ([]fsnode.T, error) {
   109  				return nil, fmt.Errorf("func was called: %q", node.Info().Name())
   110  			},
   111  		))
   112  		got := Walker{
   113  			IgnoredNames: map[string]struct{}{
   114  				addsDirName: struct{}{},
   115  			},
   116  		}.WalkContents(ctx, t, n)
   117  		assert.Equal(t, root, got)
   118  	})
   119  }