Belle II Software development
__init__.py
1#!/usr/bin/env python3
2
3
10
11"""
12b2test_utils - Helper functions useful for test scripts
13-------------------------------------------------------
14
15This module contains functions which are commonly needed for tests like changing
16log levels or switching to an empty working directory
17"""
18
19import sys
20import os
21import tempfile
22from contextlib import contextmanager
23from collections import OrderedDict
24import multiprocessing
25import basf2
26import subprocess
27import re
28from b2test_utils import logfilter
29
30
31def skip_test(reason, py_case=None):
32 """Skip a test script with a given reason. This function will end the script
33 and not return.
34
35 This is intended for scripts to be run in :ref:`b2test-scripts
36 <b2test-scripts>` and will flag the script as skipped with the given reason
37 when tests are executed.
38
39 Useful if the test depends on some external condition like a web service and
40 missing this dependency should not fail the test run.
41
42 Parameters:
43 reason (str): the reason to skip the test.
44 py_case (unittest.TestCase): if this is to be skipped within python's
45 native unittest then pass the TestCase instance
46 """
47 if py_case:
48 py_case.skipTest(reason)
49 else:
50 print(f"TEST SKIPPED: {reason}", file=sys.stderr, flush=True)
51 sys.exit(1)
52
53
54def require_file(filename, data_type="", py_case=None):
55 """Check for the existence of a test input file before attempting to open it.
56 Skips the test if not found.
57
58 Wraps `basf2.find_file` for use in test scripts run as
59 :ref`b2test-scripts <b2test-scripts>`
60
61 Parameters:
62 filename (str): relative filename to look for, either in a central place or in the current working directory
63 data_type (str): case insensitive data type to find. Either empty string or one of ``"examples"`` or ``"validation"``.
64 py_case (unittest.TestCase): if this is to be skipped within python's native unittest then pass the TestCase instance
65
66 Returns:
67 Full path to the test input file
68 """
69 try:
70 fullpath = basf2.find_file(filename, data_type, silent=False)
71 except FileNotFoundError as fnf:
72 skip_test(f'Cannot find: {fnf.filename}', py_case)
73 return fullpath
74
75
76@contextmanager
77def set_loglevel(loglevel):
78 """
79 temporarily set the log level to the specified `LogLevel <basf2.LogLevel>`. This returns a
80 context manager so it should be used in a ``with`` statement:
81
82 >>> with set_log_level(LogLevel.ERROR):
83 >>> # during this block the log level is set to ERROR
84 """
85 old_loglevel = basf2.logging.log_level
86 basf2.set_log_level(loglevel)
87 try:
88 yield
89 finally:
90 basf2.set_log_level(old_loglevel)
91
92
93@contextmanager
94def show_only_errors():
95 """temporarily set the log level to `ERROR <LogLevel.ERROR>`. This returns a
96 context manager so it should be used in a ``with`` statement
97
98 >>> with show_only_errors():
99 >>> B2INFO("this will not be shown")
100 >>> B2INFO("but this might")
101 """
102 with set_loglevel(basf2.LogLevel.ERROR):
103 yield
104
105
106def configure_logging_for_tests(user_replacements=None, replace_cdb_provider=True):
107 """
108 Change the log system to behave a bit more appropriately for testing scenarios:
109
110 1. Simplify log message to be just ``[LEVEL] message``
111 2. Disable error summary, just additional noise
112 3. Intercept all log messages and replace
113
114 * the current working directory in log messaged with ``${cwd}``
115 * the current release version with ``${release_version}``
116 * the current default globaltags with ``${default_globaltag}``
117 * the contents of the following environment variables with their name
118 (or the listed replacement string):
119
120 - :envvar:`BELLE2_TOOLS`
121 - :envvar:`BELLE2_RELEASE_DIR` with ``BELLE2_SOFTWARE_DIR``
122 - :envvar:`BELLE2_LOCAL_DIR` with ``BELLE2_SOFTWARE_DIR``
123 - :envvar:`BELLE2_EXTERNALS_DIR`
124 - :envvar:`BELLE2_VALIDATION_DATA_DIR`
125 - :envvar:`BELLE2_EXAMPLES_DATA_DIR`
126 - :envvar:`BELLE2_BACKGROUND_DIR`
127 - :envvar:`BELLE2_CONDB_METADATA`
128
129 Parameters:
130 user_replacements (dict(str, str)): Additional strings and their replacements to replace in the output
131 replace_cdb_provider (bool): If False, it does not replace the conditions database metadata provider with
132 `BELLE2_CONDB_METADATA` (necessary for some specific tests)
133
134 Warning:
135 This function should be called **after** switching directory to replace the correct directory name
136
137 .. versionadded:: release-04-00-00
138 """
139 basf2.logging.reset()
140 basf2.logging.enable_summary(False)
141 basf2.logging.enable_python_logging = True
142 basf2.logging.add_console()
143 # clang prints namespaces differently so no function names. Also let's skip the line number,
144 # we don't want failing tests just because we added a new line of code. In fact, let's just see the message
145 for level in basf2.LogLevel.values.values():
146 basf2.logging.set_info(level, basf2.LogInfo.LEVEL | basf2.LogInfo.MESSAGE)
147
148 # now create dictionary of string replacements. Since each key can only be
149 # present once order is kind of important so the less portable ones like
150 # current directory should go first and might be overridden if for example
151 # the BELLE2_LOCAL_DIR is identical to the current working directory
152 replacements = OrderedDict()
153 try:
154 replacements[basf2.version.get_version()] = "${release_version}"
155 except Exception:
156 pass
157 replacements[", ".join(basf2.conditions.default_globaltags)] = "${default_globaltag}"
158 # add a special replacement for the CDB metadata provider, since it's not set via env. variable
159 # use the first metadata provider in the list for the replacement
160 if replace_cdb_provider and len(basf2.conditions.metadata_providers) > 0:
161 replacements[basf2.conditions.metadata_providers[0]] = "${BELLE2_CONDB_METADATA}"
162 # Let's be lazy and take the environment variables from the docstring so we don't have to repeat them here
163 for env_name, replacement in re.findall(":envvar:`(.*?)`(?:.*``(.*?)``)?", configure_logging_for_tests.__doc__):
164 if not replacement:
165 replacement = env_name
166 if env_name in os.environ:
167 # replace path from the environment with the name of the variable. But remove a trailing slash or whitespace so that
168 # the output doesn't depend on whether there is a tailing slash in the environment variable
169 replacements[os.environ[env_name].rstrip('/ ')] = f"${{{replacement}}}"
170
171 if user_replacements is not None:
172 replacements.update(user_replacements)
173 # add cwd only if it doesn't overwrite anything ...
174 replacements.setdefault(os.getcwd(), "${cwd}")
175 sys.stdout = logfilter.LogReplacementFilter(sys.__stdout__, replacements)
176
177
178@contextmanager
179def working_directory(path):
180 """temporarily change the working directory to path
181
182 >>> with working_directory("testing"):
183 >>> # now in subdirectory "./testing/"
184 >>> # back to parent directory
185
186 This function will not create the directory for you. If changing into the
187 directory fails a `FileNotFoundError` will be raised.
188 """
189 dirname = os.getcwd()
190 try:
191 os.chdir(path)
192 yield
193 finally:
194 os.chdir(dirname)
195
196
197@contextmanager
198def clean_working_directory():
199 """Context manager to create a temporary directory and directly us it as
200 current working directory. The directory will automatically be deleted after
201 the with context is left.
202
203 >>> with clean_working_directory() as dirname:
204 >>> # now we are in an empty directory, name is stored in dirname
205 >>> assert(os.listdir() == [])
206 >>> # now we are back where we were before
207 """
208 with tempfile.TemporaryDirectory() as tempdir:
209 with working_directory(tempdir):
210 yield tempdir
211
212
213@contextmanager
214def local_software_directory():
215 """Context manager to make sure we are executed in the top software
216 directory by switching to $BELLE2_LOCAL_DIR or $BELLE2_RELEASE_DIR.
217
218 >>> with local_software_directory():
219 >>> assert(os.listdir().contains("analysis"))
220 """
221 directory = os.environ.get("BELLE2_LOCAL_DIR", os.environ.get("BELLE2_RELEASE_DIR", None))
222 if directory is None:
223 raise RuntimeError("Cannot find Belle II software directory, "
224 "have you setup the software correctly?")
225
226 with working_directory(directory):
227 yield directory
228
229
230def run_in_subprocess(*args, target, **kwargs):
231 """Run the given ``target`` function in a child process using `multiprocessing.Process`
232
233 This avoids side effects: anything done in the target function will not
234 affect the current process. This is mostly useful for test scripts as
235 ``target`` can emit a `FATAL <LogLevel.FATAL>` error without killing script execution.
236
237 It will return the exitcode of the child process which should be 0 in case of no error
238 """
239 process = multiprocessing.Process(target=target, args=args, kwargs=kwargs)
240 process.start()
241 process.join()
242 return process.exitcode
243
244
245def safe_process(*args, **kwargs):
246 """Run `basf2.process` with the given path in a child process using
247 `multiprocessing.Process`
248
249 This avoids side effects (`safe_process` can be safely called multiple times)
250 and doesn't kill this script even if a segmentation violation or a `FATAL
251 <LogLevel.FATAL>` error occurs during processing.
252
253 It will return the exitcode of the child process which should be 0 in case of no error
254 """
255 return run_in_subprocess(target=basf2.process, *args, **kwargs)
256
257
258def check_error_free(tool, toolname, package, filter=lambda x: False, toolopts=None):
259 """Calls the ``tool`` with argument ``package`` and check that the output is
260 error-free. Optionally ``filter`` the output in case of error messages that
261 can be ignored.
262
263 In case there is some output left, then prints the error message and exits
264 (failing the test).
265
266 Warnings:
267 If the test is skipped or the test contains errors this function does
268 not return but will directly end the program.
269
270 Arguments:
271 tool(str): executable to call
272 toolname(str): human readable name of the tool
273 package(str): package to run over. Also the first argument to the tool
274 filter: function which gets called for each line of output and
275 if it returns True the line will be ignored.
276 toolopts(list(str)): extra options to pass to the tool.
277 """
278
279 if "BELLE2_LOCAL_DIR" not in os.environ and "BELLE2_RELEASE_DIR" not in os.environ:
280 skip_test("No release is setup")
281
282 args = [tool]
283 if toolopts:
284 args += toolopts
285 if package is not None:
286 args += [package]
287
288 with local_software_directory():
289 try:
290 output = subprocess.check_output(args, encoding="utf8")
291 except subprocess.CalledProcessError as error:
292 print(error)
293 output = error.output
294
295 clean_log = [e for e in output.splitlines() if e and not filter(e)]
296 if len(clean_log) > 0:
297 subject = f"{package} package" if package is not None else "repository"
298 print(f"""\
299The {subject} has some {toolname} issues, which is now not allowed.
300Please run:
301
302 $ {" ".join(args)}
303
304and fix any issues you have introduced. Here is what {toolname} found:\n""")
305 print("\n".join(clean_log))
306 sys.exit(1)
307
308
309def get_streamer_checksums(objects):
310 """
311 Extract the version and streamer checksum of the C++ objects in the given list
312 by writing them all to a TMemFile and getting back the streamer info list
313 automatically created by ROOT afterwards.
314 Please note, that this list also includes the streamer infos of all
315 base objects of the objects you gave.
316
317 Returns a dictionary object name -> (version, checksum).
318 """
319 # Always avoid the top-level 'import ROOT'.
320 import ROOT # noqa
321
322 # Write out the objects to a mem file
323 f = ROOT.TMemFile("test_mem_file", "RECREATE")
324 f.cd()
325
326 for o in objects:
327 o.Write()
328 f.Write()
329
330 # Go through all streamer infos and extract checksum and version
331 streamer_checksums = dict()
332 for streamer_info in f.GetStreamerInfoList():
333 if not isinstance(streamer_info, ROOT.TStreamerInfo):
334 continue
335 streamer_checksums[streamer_info.GetName()] = (streamer_info.GetClassVersion(), streamer_info.GetCheckSum())
336
337 f.Close()
338 return streamer_checksums
339
340
341def get_object_with_name(object_name, root=None):
342 """
343 (Possibly) recursively get the object with the given name from the Belle2 namespace.
344
345 If the object name includes a ".", the first part will be turned into an object (probably a module)
346 and the function is continued with this object as the root and the rest of the name.
347
348 If not, the object is extracted via a getattr call.
349 """
350 if root is None:
351 from ROOT import Belle2 # noqa
352 root = Belle2
353
354 if "." in object_name:
355 namespace, object_name = object_name.split(".", 1)
356
357 return get_object_with_name(object_name, get_object_with_name(namespace, root=root))
358
359 return getattr(root, object_name)
360
361
362def skip_test_if_light(py_case=None):
363 """
364 Skips the test if we are running in a light build (maybe this test tests
365 some generation example or whatever)
366
367 Parameters:
368 py_case (unittest.TestCase): if this is to be skipped within python's
369 native unittest then pass the TestCase instance
370 """
371 try:
372 import generators # noqa
373 except ModuleNotFoundError:
374 skip_test(reason="We're in a light build.", py_case=py_case)
375
376
377def skip_test_if_central(py_case=None):
378 """
379 Skips the test if we are using a central release (and have no local
380 git repository)
381
382 Parameters:
383 py_case (unittest.TestCase): if this is to be skipped within python's
384 native unittest then pass the TestCase instance
385 """
386 if "BELLE2_RELEASE_DIR" in os.environ:
387 skip_test(reason="We're in a central release.", py_case=py_case)
388
389
390def print_belle2_environment():
391 """
392 Prints all the BELLE2 environment variables on the screen.
393 """
394 basf2.B2INFO('The BELLE2 environment variables are:')
395 for key, value in sorted(dict(os.environ).items()):
396 if 'BELLE2' in key.upper():
397 print(f' {key}={value}')
398
399
400@contextmanager
401def temporary_environment(**environ):
402 """
403 Context manager that temporarily sets environment variables
404 for the current process. Inspired by https://stackoverflow.com/a/34333710
405
406 >>> with temporary_environment(BELLE2_TEMP_DIR='/tmp/belle2'):
407 ... "BELLE2_TEMP_DIR" in os.environ
408 True
409
410 >>> "BELLE2_TEMP_DIR" in os.environ
411 False
412
413 Args:
414 **environ: Arbitrary keyword arguments specifying environment
415 variables and their values.
416 """
417 old_environ = dict(os.environ)
418 os.environ.update(environ)
419 try:
420 yield
421 finally:
422 os.environ.clear()
423 os.environ.update(old_environ)
424
425
426def is_ci() -> bool:
427 """
428 Returns true if we are running a test on our CI system (currently GitLab pipeline).
429 The 'BELLE2_IS_CI' environment variable is set on CI only when the unit
430 tests are run.
431 """
432 return os.environ.get("BELLE2_IS_CI", "no").lower() in [
433 "yes",
434 "1",
435 "y",
436 "on",
437 ]
438
439
440def is_cdb_down() -> bool:
441 """
442 Returns true if the Conditions Database (CDB) is currently unavailable or slow to respond.
443 The 'BELLE2_IS_CDB_DOWN' environment variable can be used to dynamically exclude some
444 tests that rely on the CDB in case of problems.
445 """
446 return os.environ.get("BELLE2_IS_CDB_DOWN", "no").lower() in [
447 "yes",
448 "1",
449 "y",
450 "on",
451 ]