go-hep.org/x/hep@v0.38.1/hepmc/rootcnv/rootcnv_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 rootcnv
     6  
     7  import (
     8  	"bytes"
     9  	"os"
    10  	"path/filepath"
    11  	"testing"
    12  
    13  	"go-hep.org/x/hep/groot"
    14  	"go-hep.org/x/hep/groot/riofs"
    15  	"go-hep.org/x/hep/groot/rtree"
    16  	"go-hep.org/x/hep/hepmc"
    17  	"go-hep.org/x/hep/internal/diff"
    18  )
    19  
    20  func TestRW(t *testing.T) {
    21  	dir, err := os.MkdirTemp("", "hepmc-rootcnv-")
    22  	if err != nil {
    23  		t.Fatal(err)
    24  	}
    25  	defer os.RemoveAll(dir)
    26  
    27  	for _, tc := range []string{
    28  		"../testdata/small.hepmc",
    29  		"../testdata/test.hepmc",
    30  	} {
    31  		t.Run(tc, func(t *testing.T) {
    32  			raw, err := os.ReadFile(tc)
    33  			if err != nil {
    34  				t.Fatal(err)
    35  			}
    36  
    37  			fname := filepath.Join(dir, filepath.Base(tc)+".root")
    38  			o, err := groot.Create(fname)
    39  			if err != nil {
    40  				t.Fatalf("could not create output ROOT file: %+v", err)
    41  			}
    42  			defer o.Close()
    43  
    44  			w, err := NewFlatTreeWriter(o, "tree", rtree.WithTitle("HepMC ROOT tree"))
    45  			if err != nil {
    46  				t.Fatalf("could not create ROOT tree writer: %+v", err)
    47  			}
    48  
    49  			_, err = hepmc.Copy(w, hepmc.NewASCIIReader(bytes.NewReader(raw)))
    50  			if err != nil {
    51  				t.Fatalf("could not copy hepmc event to ROOT: %+v", err)
    52  			}
    53  
    54  			err = w.Close()
    55  			if err != nil {
    56  				t.Fatalf("could not close ROOT tree writer: %+v", err)
    57  			}
    58  
    59  			err = o.Close()
    60  			if err != nil {
    61  				t.Fatalf("could not close ROOT file: %+v", err)
    62  			}
    63  
    64  			f, err := groot.Open(fname)
    65  			if err != nil {
    66  				t.Fatalf("could not open ROOT file: %+v", err)
    67  			}
    68  			defer f.Close()
    69  
    70  			r, err := riofs.Get[rtree.Tree](f, "tree")
    71  			if err != nil {
    72  				t.Fatalf("could not retrieve ROOT tree: %+v", err)
    73  			}
    74  
    75  			rr, err := NewFlatTreeReader(r)
    76  			if err != nil {
    77  				t.Fatalf("could not create ROOT tree reader: %+v", err)
    78  			}
    79  			defer rr.Close()
    80  
    81  			buf := new(bytes.Buffer)
    82  
    83  			ww := hepmc.NewASCIIWriter(buf)
    84  			_, err = hepmc.Copy(ww, rr)
    85  			if err != nil {
    86  				t.Fatalf("could not copy hepmc event from ROOT: %+v", err)
    87  			}
    88  
    89  			err = ww.Close()
    90  			if err != nil {
    91  				t.Fatalf("could not close hepmc writer: %+v", err)
    92  			}
    93  
    94  			if got, want := buf.String(), string(raw); got != want {
    95  				d := diff.Format(string(got), string(want))
    96  				t.Fatalf("invalid r/w round trip:\n%s", d)
    97  			}
    98  		})
    99  	}
   100  }