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