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