tests/helpers/__init__.py

0001 import os
0002 import shutil
0003 from typing import List, Union
0004 import contextlib
0005
0006 import acts
0007 from acts.examples import IAlgorithm
0008
0009 geant4Enabled = (
0010     any(v.startswith("G4") for v in os.environ.keys())
0011     or "GEANT4_DATA_DIR" in os.environ
0012 )
0013 if geant4Enabled:
0014     try:
0015         import acts.examples.geant4
0016     except ImportError:
0017         geant4Enabled = False
0018
0019 try:
0020     import ROOT
0021
0022     rootEnabled = True
0023 except ImportError:
0024     rootEnabled = False
0025
0026     if "ROOTSYS" in os.environ:  # ROOT seems to be set up, but no PyROOT
0027         import warnings
0028
0029         warnings.warn(
0030             "ROOT likely built without/with incompatible PyROOT. Skipping tests that need ROOT"
0031         )
0032
0033 dd4hepEnabled = "DD4hep_DIR" in os.environ
0034 if dd4hepEnabled:
0035     try:
0036         import acts.examples.dd4hep
0037     except ImportError:
0038         dd4hepEnabled = False
0039
0040 try:
0041     import acts.examples.hepmc3
0042
0043     hepmc3Enabled = True
0044 except ImportError:
0045     hepmc3Enabled = False
0046
0047 try:
0048     import acts.examples.edm4hep
0049
0050     edm4hepEnabled = True
0051 except ImportError:
0052     edm4hepEnabled = False
0053
0054 try:
0055     import acts.examples.onnx
0056
0057     onnxEnabled = True
0058 except ImportError:
0059     onnxEnabled = False
0060
0061 try:
0062     from acts import covfie
0063
0064     covfieEnabled = True
0065 except ImportError:
0066     covfieEnabled = False
0067
0068
0069 try:
0070     import acts.examples
0071
0072     pythia8Enabled = hasattr(acts.examples, "pythia8")
0073 except ImportError:
0074     pythia8Enabled = False
0075
0076 try:
0077     import acts.examples
0078
0079     hashingSeedingEnabled = hasattr(acts.examples, "hashing")
0080 except ImportError:
0081     hashingSeedingEnabled = False
0082
0083
0084 exatrkxEnabled = shutil.which("nvidia-smi") is not None
0085 if exatrkxEnabled:
0086     try:
0087         from acts.examples import TrackFindingAlgorithmExaTrkX
0088     except ImportError:
0089         exatrkxEnabled = False
0090
0091 try:
0092     import podio
0093
0094     podioEnabled = True
0095 except ModuleNotFoundError:
0096     podioEnabled = False
0097 except ImportError:
0098     podioEnabled = False
0099
0100 isCI = os.environ.get("CI") is not None
0101
0102 if isCI:
0103     for k, v in dict(locals()).items():
0104         if k.endswith("Enabled"):
0105             locals()[k] = True
0106
0107
0108 class AssertCollectionExistsAlg(IAlgorithm):
0109     events_seen = 0
0110     collections: List[str]
0111
0112     def __init__(
0113         self,
0114         collections: Union[List[str], str],
0115         name="check_alg",
0116         level=acts.logging.INFO,
0117         *args,
0118         **kwargs,
0119     ):
0120         if isinstance(collections, str):
0121             self.collections = [collections]
0122         else:
0123             self.collections = collections
0124         IAlgorithm.__init__(self, name=name, level=level, *args, **kwargs)
0125
0126     def execute(self, ctx):
0127         try:
0128             for collection in self.collections:
0129                 assert ctx.eventStore.exists(collection), f"{collection} does not exist"
0130             self.events_seen += 1
0131             return acts.examples.ProcessCode.SUCCESS
0132         except AssertionError:
0133             print("Available collections:")
0134             print(ctx.eventStore.keys)
0135             raise
0136
0137
0138 doHashChecks = os.environ.get("ROOT_HASH_CHECKS", "") != "" or "CI" in os.environ
0139
0140
0141 @contextlib.contextmanager
0142 def failure_threshold(level: acts.logging.Level, enabled: bool = True):
0143     prev = acts.logging.getFailureThreshold()
0144     if enabled and prev != level:
0145         try:
0146             acts.logging.setFailureThreshold(level)
0147         except RuntimeError:
0148             # Repackage with different error string
0149             raise RuntimeError(
0150                 "Runtime log failure threshold could not be set. "
0151                 "Compile-time value is probably set via CMake, i.e. "
0152                 f"`ACTS_LOG_FAILURE_THRESHOLD={acts.logging.getFailureThreshold().name}` is set, "
0153                 "or `ACTS_ENABLE_LOG_FAILURE_THRESHOLD=OFF`. "
0154                 "The pytest test-suite will not work in this configuration."
0155             )
0156
0157         yield
0158         acts.logging.setFailureThreshold(prev)
0159     else:
0160         yield