Belle II Software light-2406-ragdoll
mcevent_forwarding.py
1#!/usr/bin/env python3
2
3
10
11import sys
12import basf2
13from b2test_utils import clean_working_directory, safe_process, show_only_errors
14from ROOT import Belle2
15import subprocess
16import json
17
18
19def check_mc_events(outfile, expect):
20 """Check if the number of MC events reported by the file metadata is what whe expect"""
21 try:
22 metadata = subprocess.check_output(["b2file-metadata-show", "--json", outfile])
23 mcEvents = json.loads(metadata)["mcEvents"]
24 except Exception as e:
25 basf2.B2ERROR(f"Cannot determine number of MC events: {e}")
26 return False
27
28 if mcEvents != expect:
29 basf2.B2ERROR(f"{outfile} MCEvents not correct: got {mcEvents}, expected {expect}")
30 return False
31 else:
32 basf2.B2INFO(f"{outfile}: MCEvents=={expect} Ok")
33 return True
34
35
36def run_rootio(outfile, infiles, expect, *, N=0, **input_kwargs):
37 """
38 Run RootInput/RootOutput and check that the number of events is what we want
39
40 Arguments:
41 outfile (str): filename for the output
42 infiles (list): list of input filenames
43 expect (int): numer of MC events we expect in the final file
44 N (int): pass to process() as second argument to limit the number of events
45
46 All other arguments are passed to RootInput as module parameters
47 """
48
49 global failures
50 path = basf2.create_path()
51 path.add_module("RootInput", inputFileNames=infiles, **input_kwargs)
52 path.add_module("RootOutput", outputFileName=outfile)
53 # zero counters so we can continue processing even if we already reported
54 # errors before
55 basf2.logging.zero_counters()
56 with show_only_errors():
57 safe_process(path, N)
58
59 result = check_mc_events(outfile, expect)
60 failures += 0 if result else 1
61 return result
62
63
64# no summary and load dictionaries early
65basf2.logging.enable_summary(False)
66Belle2.Environment.Instance().setNumberEventsOverride(0)
67
68# so let's go in an empty directory
69with clean_working_directory():
70 # and create a few files with some empty events
71 all_files = []
72 event_counts = []
73 failures = 0
74 for i in range(1, 16, 7):
75 all_files.append(f"events-{i:03d}.root")
76 event_counts.append(i)
77 generate = basf2.create_path()
78 generate.add_module("EventInfoSetter", evtNumList=i)
79 generate.add_module("RootOutput", outputFileName=all_files[-1])
80 with show_only_errors():
81 safe_process(generate)
82
83 # now for either one or all of those files run a number of checks:
84 for name, total, filenames, sequence in [
85 ("single", event_counts[-1], all_files[-1:], ["1:3"]),
86 ("total", sum(event_counts), all_files, ["0", "1:2", "2-"]),
87 ]:
88 # simple processing should forward MC events
89 run_rootio(f"{name}.root", filenames, total)
90 # but any skipping should not
91 run_rootio(f"{name}-skip1.root", filenames, 0, skipNEvents=1)
92 run_rootio(f"{name}-skipTo.root", filenames, 0, skipToEvent=[0, 0, 1])
93 run_rootio(f"{name}-sequence.root", filenames, 0, entrySequences=sequence)
94 # and limiting the number of events should only forward if all events
95 # are processed
96 for N in [1, total-1, total, 100000]:
97 Belle2.Environment.Instance().setNumberEventsOverride(N)
98 run_rootio(f"{name}-basf2-n-{N}.root", filenames, 0 if N < total else total)
99 Belle2.Environment.Instance().setNumberEventsOverride(0)
100 run_rootio(f"{name}-proc-{N}.root", filenames, 0 if N < total else total, N=N)
101
102 # and if some files are processed multiple times we don't want to forward
103 # the number either
104 run_rootio("duplicate.root", all_files*2, 0)
105
106if failures > 0:
107 basf2.B2ERROR(f"{failures} tests failed ...")
108 sys.exit(1)
static Environment & Instance()
Static method to get a reference to the Environment instance.
Definition: Environment.cc:28