go-hep.org/x/hep@v0.38.1/groot/rnpy/column_test.go (about)

     1  // Copyright ©2022 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 rnpy
     6  
     7  import (
     8  	"fmt"
     9  	"reflect"
    10  	"testing"
    11  
    12  	"go-hep.org/x/hep/groot/riofs"
    13  	"go-hep.org/x/hep/groot/rtree"
    14  )
    15  
    16  func TestColumn(t *testing.T) {
    17  	for _, tc := range []struct {
    18  		fname string
    19  		tname string
    20  		aname rtree.ReadVar
    21  		want  any
    22  		err   error
    23  	}{
    24  		{
    25  			fname: "../testdata/simple.root",
    26  			tname: "tree",
    27  			aname: rtree.ReadVar{Name: "one"},
    28  			want:  []int32{1, 2, 3, 4},
    29  		},
    30  		{
    31  			fname: "../testdata/simple.root",
    32  			tname: "tree",
    33  			aname: rtree.ReadVar{Name: "two"},
    34  			want:  []float32{1.1, 2.2, 3.3, 4.4},
    35  		},
    36  		{
    37  			fname: "../testdata/simple.root",
    38  			tname: "tree",
    39  			aname: rtree.ReadVar{Name: "three"},
    40  			want:  []string{"uno", "dos", "tres", "quatro"},
    41  		},
    42  		{
    43  			fname: "../testdata/leaves.root",
    44  			tname: "tree",
    45  			aname: rtree.ReadVar{Name: "ArrF64"},
    46  			want: func() any {
    47  				o := make([][10]float64, 10)
    48  				for i := range o {
    49  					for j := range o[i] {
    50  						o[i][j] = float64(i)
    51  					}
    52  				}
    53  				return o
    54  			}(),
    55  		},
    56  		{
    57  			fname: "../testdata/root_numpy_struct.root",
    58  			tname: "test",
    59  			aname: rtree.ReadVar{Name: "branch1", Leaf: "intleaf"},
    60  			want:  []int32{10},
    61  		},
    62  		{
    63  			fname: "../testdata/root_numpy_struct.root",
    64  			tname: "test",
    65  			aname: rtree.ReadVar{Name: "branch1", Leaf: "floatleaf"},
    66  			want:  []float32{15.5},
    67  		},
    68  		{
    69  			fname: "../testdata/root_numpy_struct.root",
    70  			tname: "test",
    71  			aname: rtree.ReadVar{Name: "branch2", Leaf: "intleaf"},
    72  			want:  []int32{20},
    73  		},
    74  		{
    75  			fname: "../testdata/root_numpy_struct.root",
    76  			tname: "test",
    77  			aname: rtree.ReadVar{Name: "branch2", Leaf: "floatleaf"},
    78  			want:  []float32{781.2},
    79  		},
    80  		{
    81  			fname: "../testdata/simple.root",
    82  			tname: "tree",
    83  			aname: rtree.ReadVar{Name: "not_there"},
    84  			err:   fmt.Errorf(`rnpy: no rvar named "not_there"`),
    85  		},
    86  		{
    87  			fname: "../testdata/root_numpy_struct.root",
    88  			tname: "test",
    89  			aname: rtree.ReadVar{Name: "branch1", Leaf: "not_there"},
    90  			err:   fmt.Errorf(`rnpy: no rvar named "branch1.not_there"`),
    91  		},
    92  		{
    93  			fname: "../testdata/root_numpy_struct.root",
    94  			tname: "test",
    95  			aname: rtree.ReadVar{Name: "branch2", Leaf: "not_there"},
    96  			err:   fmt.Errorf(`rnpy: no rvar named "branch2.not_there"`),
    97  		},
    98  		{
    99  			fname: "../testdata/leaves.root",
   100  			tname: "tree",
   101  			aname: rtree.ReadVar{Name: "SliF64"},
   102  			err:   fmt.Errorf(`rnpy: invalid branch or leaf type []float64`),
   103  		},
   104  	} {
   105  		t.Run("", func(t *testing.T) {
   106  			f, err := riofs.Open(tc.fname)
   107  			if err != nil {
   108  				t.Fatal(err)
   109  			}
   110  			defer f.Close()
   111  
   112  			obj, err := riofs.Dir(f).Get(tc.tname)
   113  			if err != nil {
   114  				t.Fatal(err)
   115  			}
   116  			tree := obj.(rtree.Tree)
   117  
   118  			var sli any
   119  
   120  			col, err := NewColumn(tree, tc.aname)
   121  			if err == nil {
   122  				sli, err = col.Slice()
   123  			}
   124  
   125  			switch {
   126  			case err != nil && tc.err != nil:
   127  				if got, want := err.Error(), tc.err.Error(); got != want {
   128  					t.Fatalf("invalid error:\ngot= %s\nwant=%s", got, want)
   129  				}
   130  				return
   131  			case err != nil && tc.err == nil:
   132  				t.Fatalf("could not create column: %+v", err)
   133  			case err == nil && tc.err != nil:
   134  				t.Fatalf("expected an error: %+v", tc.err)
   135  			case err == nil && tc.err == nil:
   136  				// ok.
   137  			}
   138  
   139  			if got, want := sli, tc.want; !reflect.DeepEqual(got, want) {
   140  				t.Fatalf("invalid slice:\ngot= %+v\nwant=%+v", got, want)
   141  			}
   142  		})
   143  	}
   144  }
   145  
   146  func TestColumnReadFail(t *testing.T) {
   147  	f, err := riofs.Open("../testdata/leaves.root")
   148  	if err != nil {
   149  		t.Fatal(err)
   150  	}
   151  	defer f.Close()
   152  
   153  	obj, err := f.Get("tree")
   154  	if err != nil {
   155  		t.Fatal(err)
   156  	}
   157  	tree := obj.(rtree.Tree)
   158  
   159  	col, err := NewColumn(tree, rtree.ReadVar{Name: "F64"})
   160  	if err != nil {
   161  		t.Fatal(err)
   162  	}
   163  
   164  	want := fmt.Errorf(`rnpy: could not create ROOT reader for "boo": rtree: could not create reader: rtree: tree "tree" has no branch named "boo"`)
   165  
   166  	col.rvar.Name = "boo"
   167  	_, err = col.Slice()
   168  	if got, want := err.Error(), want.Error(); got != want {
   169  		t.Fatalf("invalid error:\ngot= %+v\nwant=%+v", got, want)
   170  	}
   171  }