go-hep.org/x/hep@v0.38.1/groot/rcmd/split_test.go (about)

     1  // Copyright ©2020 The go-hep Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package rcmd_test
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"os"
    11  	"path/filepath"
    12  	"testing"
    13  
    14  	"go-hep.org/x/hep/groot"
    15  	"go-hep.org/x/hep/groot/rcmd"
    16  	"go-hep.org/x/hep/groot/riofs"
    17  	"go-hep.org/x/hep/groot/rtree"
    18  )
    19  
    20  func TestSplit(t *testing.T) {
    21  	tmp, err := os.MkdirTemp("", "groot-root-split-")
    22  	if err != nil {
    23  		t.Fatalf("%+v", err)
    24  	}
    25  	defer os.RemoveAll(tmp)
    26  
    27  	type funcT func(t *testing.T, fname string) error
    28  	for _, tc := range []struct {
    29  		name    string
    30  		n       int64
    31  		input   funcT
    32  		outputs []funcT
    33  	}{
    34  		{
    35  			name:  "flat-tree-1",
    36  			n:     10,
    37  			input: makeSplitFlatTree(0, 10),
    38  			outputs: []funcT{
    39  				makeSplitFlatTree(0, 10),
    40  			},
    41  		},
    42  		{
    43  			name:  "flat-tree-2",
    44  			n:     5,
    45  			input: makeSplitFlatTree(0, 10),
    46  			outputs: []funcT{
    47  				makeSplitFlatTree(0, 5),
    48  				makeSplitFlatTree(5, 10),
    49  			},
    50  		},
    51  		{
    52  			name:  "flat-tree-3",
    53  			n:     11,
    54  			input: makeSplitFlatTree(0, 10),
    55  			outputs: []funcT{
    56  				makeSplitFlatTree(0, 10),
    57  			},
    58  		},
    59  	} {
    60  		t.Run(tc.name, func(t *testing.T) {
    61  			var (
    62  				fname   = filepath.Join(tmp, tc.name+".in.root")
    63  				oname   = filepath.Join(tmp, tc.name+".out.root")
    64  				deep    = true
    65  				verbose = true
    66  			)
    67  
    68  			err := tc.input(t, fname)
    69  			if err != nil {
    70  				t.Fatalf("%+v", err)
    71  			}
    72  
    73  			fnames, err := rcmd.Split(oname, fname, "dir-1/dir-11/mytree", tc.n, verbose)
    74  			if err != nil {
    75  				t.Fatalf("could not run root-merge: %+v", err)
    76  			}
    77  			if got, want := len(fnames), len(tc.outputs); got != want {
    78  				t.Fatalf("invalid number of split files: got=%d, want=%d", got, want)
    79  			}
    80  
    81  			for i, wantFunc := range tc.outputs {
    82  				oname := filepath.Join(tmp, fmt.Sprintf(tc.name+".out-%d.root", i))
    83  				got := new(bytes.Buffer)
    84  				err = rcmd.Dump(got, oname, deep, nil)
    85  				if err != nil {
    86  					t.Fatalf("could not run root-dump: %+v", err)
    87  				}
    88  
    89  				refname := filepath.Join(tmp, fmt.Sprintf(tc.name+"-%d.want.root", i))
    90  				err = wantFunc(t, refname)
    91  				if err != nil {
    92  					t.Fatalf("%+v", err)
    93  				}
    94  				want := new(bytes.Buffer)
    95  				err = rcmd.Dump(want, refname, deep, nil)
    96  				if err != nil {
    97  					t.Fatalf("could not run root-dump: %+v", err)
    98  				}
    99  
   100  				if got, want := got.String(), want.String(); got != want {
   101  					t.Fatalf("invalid root-merge output:\ngot:\n%swant:\n%s", got, want)
   102  				}
   103  			}
   104  		})
   105  	}
   106  }
   107  
   108  func makeSplitFlatTree(beg, end int) func(t *testing.T, fname string) error {
   109  	return func(t *testing.T, fname string) error {
   110  		type Data struct {
   111  			I32    int32
   112  			F64    float64
   113  			Str    string
   114  			ArrF64 [5]float64
   115  			N      int32
   116  			SliF64 []float64 `groot:"SliF64[N]"`
   117  		}
   118  
   119  		f, err := groot.Create(fname)
   120  		if err != nil {
   121  			t.Fatalf("%+v", err)
   122  		}
   123  		defer f.Close()
   124  
   125  		dir, err := riofs.Dir(f).Mkdir("dir-1/dir-11")
   126  		if err != nil {
   127  			t.Fatalf("could not create directory: %+v", err)
   128  		}
   129  
   130  		var evt Data
   131  		tree, err := rtree.NewWriter(dir, "mytree", rtree.WriteVarsFromStruct(&evt))
   132  		if err != nil {
   133  			t.Fatalf("could not create tree writer: %+v", err)
   134  		}
   135  
   136  		for i := beg; i < end; i++ {
   137  			evt.I32 = int32(i)
   138  			evt.F64 = float64(i)
   139  			evt.Str = fmt.Sprintf("evt-%0d", i)
   140  			evt.ArrF64 = [5]float64{float64(i), float64(i + 1), float64(i + 2), float64(i + 3), float64(i + 4)}
   141  			j := i % 5
   142  			evt.N = int32(j)
   143  			evt.SliF64 = []float64{float64(i), float64(i + 1), float64(i + 2), float64(i + 3), float64(i + 4)}[:j]
   144  			_, err = tree.Write()
   145  			if err != nil {
   146  				t.Fatalf("could not write event %d: %+v", i, err)
   147  			}
   148  		}
   149  
   150  		err = tree.Close()
   151  		if err != nil {
   152  			t.Fatalf("could not write tree: %+v", err)
   153  		}
   154  
   155  		err = f.Close()
   156  		if err != nil {
   157  			t.Fatalf("could not close file: %+v", err)
   158  		}
   159  
   160  		return nil
   161  	}
   162  }