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 try:
0034     import acts
0035
0036     geomodelEnabled = hasattr(acts, "geomodel")
0037 except ImportError:
0038     geomodelEnabled = False
0039
0040 dd4hepEnabled = "DD4hep_DIR" in os.environ
0041 if dd4hepEnabled:
0042     try:
0043         import acts.examples.dd4hep
0044     except ImportError:
0045         dd4hepEnabled = False
0046
0047 try:
0048     import acts.examples.hepmc3
0049
0050     hepmc3Enabled = True
0051 except ImportError:
0052     hepmc3Enabled = False
0053
0054 try:
0055     import acts.examples.edm4hep
0056
0057     edm4hepEnabled = True
0058 except ImportError:
0059     edm4hepEnabled = False
0060
0061 try:
0062     import acts.examples
0063
0064     onnxEnabled = hasattr(acts.examples, "onnx")
0065 except ImportError:
0066     onnxEnabled = False
0067
0068 try:
0069     from acts import covfie
0070
0071     covfieEnabled = True
0072 except ImportError:
0073     covfieEnabled = False
0074
0075
0076 try:
0077     import acts.examples
0078
0079     pythia8Enabled = hasattr(acts.examples, "pythia8")
0080 except ImportError:
0081     pythia8Enabled = False
0082
0083 try:
0084     import acts.examples
0085
0086     hashingSeedingEnabled = hasattr(acts.examples, "hashing")
0087 except ImportError:
0088     hashingSeedingEnabled = False
0089
0090
0091 gnnEnabled = shutil.which("nvidia-smi") is not None
0092 if gnnEnabled:
0093     try:
0094         from acts.examples import TrackFindingAlgorithmGnn
0095     except ImportError:
0096         gnnEnabled = False
0097
0098 try:
0099     import podio
0100
0101     podioEnabled = True
0102 except ModuleNotFoundError:
0103     podioEnabled = False
0104 except ImportError:
0105     podioEnabled = False
0106
0107 isCI = os.environ.get("CI") is not None
0108
0109
0110 class AssertCollectionExistsAlg(IAlgorithm):
0111     events_seen = 0
0112     collections: List[str]
0113
0114     def __init__(
0115         self,
0116         collections: Union[List[str], str],
0117         name="check_alg",
0118         level=acts.logging.INFO,
0119         *args,
0120         **kwargs,
0121     ):
0122         if isinstance(collections, str):
0123             self.collections = [collections]
0124         else:
0125             self.collections = collections
0126         IAlgorithm.__init__(self, name=name, level=level, *args, **kwargs)
0127
0128     def execute(self, ctx):
0129         try:
0130             for collection in self.collections:
0131                 assert ctx.eventStore.exists(collection), f"{collection} does not exist"
0132             self.events_seen += 1
0133             return acts.examples.ProcessCode.SUCCESS
0134         except AssertionError:
0135             print("Available collections:")
0136             print(ctx.eventStore.keys)
0137             raise
0138
0139
0140 doHashChecks = False
0141 _hashEnvVar = os.environ.get("ROOT_HASH_CHECKS")
0142
0143 if _hashEnvVar is not None:
0144     if _hashEnvVar.lower() not in ("off", "0", "false"):
0145         doHashChecks = True
0146 else:
0147     if "CI" in os.environ:
0148         doHashChecks = True
0149
0150
0151 @contextlib.contextmanager
0152 def failure_threshold(level: acts.logging.Level, enabled: bool = True):
0153     prev = acts.logging.getFailureThreshold()
0154     if enabled and prev != level:
0155         try:
0156             acts.logging.setFailureThreshold(level)
0157         except RuntimeError:
0158             # Repackage with different error string
0159             raise RuntimeError(
0160                 "Runtime log failure threshold could not be set. "
0161                 "Compile-time value is probably set via CMake, i.e. "
0162                 f"`ACTS_LOG_FAILURE_THRESHOLD={acts.logging.getFailureThreshold().name}` is set, "
0163                 "or `ACTS_ENABLE_LOG_FAILURE_THRESHOLD=OFF`. "
0164                 "The pytest test-suite will not work in this configuration."
0165             )
0166
0167         yield
0168         acts.logging.setFailureThreshold(prev)
0169     else:
0170         yield