Belle II Software light-2406-ragdoll
include_exclude_branches.py
1
8import basf2 as b2
9import b2test_utils as b2u
10import ROOT
11
12
13def filter_input_output(input_file_name, output_file_name, include_branches=None, exclude_branches=None):
14
15 b2u.configure_logging_for_tests()
16 b2.set_random_seed('Forza Livorno!!!')
17 b2.B2INFO('We are testing the following configuration:',
18 include_branches=include_branches,
19 exclude_branches=exclude_branches)
20 child_path = b2.Path()
21 child_path.add_module('RootInput',
22 inputFileName=input_file_name,
23 collectStatistics=True)
24 if include_branches is not None:
25 b2.set_module_parameters(child_path,
26 'RootInput',
27 branchNames=include_branches)
28 if exclude_branches is not None:
29 b2.set_module_parameters(child_path,
30 'RootInput',
31 excludeBranchNames=exclude_branches)
32 child_path.add_module('RootOutput',
33 outputFileName=output_file_name)
34 b2.process(child_path)
35
36
37def check_file_content(file_name, input_branches):
38
39 f = ROOT.TFile.Open(file_name, 'READ')
40 if f.IsZombie():
41 b2.B2FATAL(f'Cannot open the input file {file_name}.')
42 tree = f.Get('tree')
43 file_branches = [branch.GetName() for branch in tree.GetListOfBranches()]
44 f.Close()
45 return set(input_branches) <= set(file_branches)
46
47
48if __name__ == '__main__':
49
50 with b2u.clean_working_directory():
51 # Let's check the input file
52 input_file_name = b2.find_file('framework/tests/mdst_with_muid.root')
53 assert (0 == b2u.run_in_subprocess(input_file_name, 'mdst_all_branches.root', None, None,
54 target=filter_input_output))
55 all_branches = ['Tracks', 'ECLClusters', 'TracksToECLClusters', 'Muids', 'TracksToMuids']
56 assert (1 == check_file_content(input_file_name, all_branches))
57 # Let's exclude Muids.
58 file_to_check = 'mdst_without_muid.root'
59 assert (0 == b2u.run_in_subprocess(input_file_name, file_to_check, None, ['Muids'],
60 target=filter_input_output))
61 assert (0 == check_file_content(file_to_check, ['Muids']))
62 assert (0 == check_file_content(file_to_check, ['TracksToMuids']))
63 assert (1 == check_file_content(file_to_check, ['Tracks', 'ECLClusters']))
64 # Let's include only Muids.
65 file_to_check = 'mdst_only_muid.root'
66 assert (0 == b2u.run_in_subprocess(input_file_name, file_to_check, ['Muids'], None,
67 target=filter_input_output))
68 # The dictionary is not found, so it can't be stored.
69 assert (0 == check_file_content(file_to_check, ['Muids']))
70 # Let's include only Tracks and ECLClusters.
71 assert (0 == b2u.run_in_subprocess(input_file_name, file_to_check, ['Tracks', 'ECLClusters'], None,
72 target=filter_input_output))
73 assert (0 == check_file_content(file_to_check, ['Muids', 'TracksToMuids']))
74 assert (1 == check_file_content(file_to_check, ['Tracks', 'ECLClusters', 'TracksToECLClusters']))