PhysicsTools/HeppyCore/python/statistics/tree_test.py - master - githubmirror/cmssw

PhysicsTools/HeppyCore/python/statistics/tree_test.py

35 строк804 B

Shahzad Malik Muzaffar

[ANALYSIS] py2/3 compatibility:drop use of __future__

22 ноя 2024, 20:16

22 ноя 2024, 20:168d027cc

import unittest
from ROOT import TFile
from .tree import Tree

class TreeTestCase(unittest.TestCase):

    def test_fill(self):
        fi = TFile('tree.root','RECREATE')
        tr = Tree('test_tree', 'A test tree')
        tr.var('a')
        tr.var('b')
        tr.fill('a', 3)
        tr.fill('a', 4)
        tr.fill('b', 5)
        tr.tree.Fill()
        fi.Write()
        fi.Close()

    def test_read(self):
        fi = TFile('tree.root')
        tr = fi.Get('test_tree')
        self.assertEqual(tr.GetEntries(), 1)
        tr.GetEntry(0)
        self.assertEqual(tr.a, 4)

    def test_iterate(self):
        fi = TFile('tree.root')
        tr = fi.Get('test_tree')
        for ev in tr:
            self.assertEqual(ev.a, 4)
            break


if __name__ == '__main__':
    unittest.main()