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
0062 try:
0063     import acts.examples
0064
0065     pythia8Enabled = hasattr(acts.examples, "pythia8")
0066 except ImportError:
0067     pythia8Enabled = False
0068
0069
0070 exatrkxEnabled = shutil.which("nvidia-smi") is not None
0071 if exatrkxEnabled:
0072     try:
0073         from acts.examples import TrackFindingAlgorithmExaTrkX
0074     except ImportError:
0075         exatrkxEnabled = False
0076
0077 try:
0078     import podio
0079
0080     podioEnabled = True
0081 except ModuleNotFoundError:
0082     podioEnabled = False
0083
0084 isCI = os.environ.get("CI", "false") == "true"
0085
0086 if isCI:
0087     for k, v in dict(locals()).items():
0088         if k.endswith("Enabled"):
0089             locals()[k] = True
0090
0091
0092 class AssertCollectionExistsAlg(IAlgorithm):
0093     events_seen = 0
0094     collections: List[str]
0095
0096     def __init__(
0097         self,
0098         collections: Union[List[str], str],
0099         name="check_alg",
0100         level=acts.logging.INFO,
0101         *args,
0102         **kwargs,
0103     ):
0104         if isinstance(collections, str):
0105             self.collections = [collections]
0106         else:
0107             self.collections = collections
0108         IAlgorithm.__init__(self, name=name, level=level, *args, **kwargs)
0109
0110     def execute(self, ctx):
0111         for collection in self.collections:
0112             assert ctx.eventStore.exists(collection), f"{collection} does not exist"
0113         self.events_seen += 1
0114         return acts.examples.ProcessCode.SUCCESS
0115
0116
0117 doHashChecks = os.environ.get("ROOT_HASH_CHECKS", "") != "" or "CI" in os.environ
0118
0119
0120 @contextlib.contextmanager
0121 def failure_threshold(level: acts.logging.Level, enabled: bool = True):
0122     prev = acts.logging.getFailureThreshold()
0123     if enabled and prev != level:
0124         try:
0125             acts.logging.setFailureThreshold(level)
0126         except RuntimeError:
0127             # Repackage with different error string
0128             raise RuntimeError(
0129                 "Runtime log failure threshold could not be set. "
0130                 "Compile-time value is probably set via CMake, i.e. "
0131                 f"`ACTS_LOG_FAILURE_THRESHOLD={acts.logging.getFailureThreshold().name}` is set, "
0132                 "or `ACTS_ENABLE_LOG_FAILURE_THRESHOLD=OFF`. "
0133                 "The pytest test-suite will not work in this configuration."
0134             )
0135
0136         yield
0137         acts.logging.setFailureThreshold(prev)
0138     else:
0139         yield