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.onnx
0063
0064     onnxEnabled = True
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 if isCI:
0110     for k, v in dict(locals()).items():
0111         if k.endswith("Enabled"):
0112             locals()[k] = True
0113
0114
0115 class AssertCollectionExistsAlg(IAlgorithm):
0116     events_seen = 0
0117     collections: List[str]
0118
0119     def __init__(
0120         self,
0121         collections: Union[List[str], str],
0122         name="check_alg",
0123         level=acts.logging.INFO,
0124         *args,
0125         **kwargs,
0126     ):
0127         if isinstance(collections, str):
0128             self.collections = [collections]
0129         else:
0130             self.collections = collections
0131         IAlgorithm.__init__(self, name=name, level=level, *args, **kwargs)
0132
0133     def execute(self, ctx):
0134         try:
0135             for collection in self.collections:
0136                 assert ctx.eventStore.exists(collection), f"{collection} does not exist"
0137             self.events_seen += 1
0138             return acts.examples.ProcessCode.SUCCESS
0139         except AssertionError:
0140             print("Available collections:")
0141             print(ctx.eventStore.keys)
0142             raise
0143
0144
0145 doHashChecks = os.environ.get("ROOT_HASH_CHECKS", "") != "" or "CI" in os.environ
0146
0147
0148 @contextlib.contextmanager
0149 def failure_threshold(level: acts.logging.Level, enabled: bool = True):
0150     prev = acts.logging.getFailureThreshold()
0151     if enabled and prev != level:
0152         try:
0153             acts.logging.setFailureThreshold(level)
0154         except RuntimeError:
0155             # Repackage with different error string
0156             raise RuntimeError(
0157                 "Runtime log failure threshold could not be set. "
0158                 "Compile-time value is probably set via CMake, i.e. "
0159                 f"`ACTS_LOG_FAILURE_THRESHOLD={acts.logging.getFailureThreshold().name}` is set, "
0160                 "or `ACTS_ENABLE_LOG_FAILURE_THRESHOLD=OFF`. "
0161                 "The pytest test-suite will not work in this configuration."
0162             )
0163
0164         yield
0165         acts.logging.setFailureThreshold(prev)
0166     else:
0167         yield