From bfed86eb687abb40018bb9e54cdf775b64bb7e26 Mon Sep 17 00:00:00 2001 From: Ara Ghukasyan Date: Thu, 18 Jan 2024 11:41:40 -0500 Subject: [PATCH] wrap import covalent for convenience --- covalent_slurm_plugin/exec.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/covalent_slurm_plugin/exec.py b/covalent_slurm_plugin/exec.py index 9ca1f5d..4ccf366 100644 --- a/covalent_slurm_plugin/exec.py +++ b/covalent_slurm_plugin/exec.py @@ -22,6 +22,11 @@ import cloudpickle as pickle +def _import_covalent() -> None: + # Wrapped import for convenience in testing. + import covalent + + def _check_setup() -> None: """Use these checks to create more informative error messages.""" @@ -33,7 +38,7 @@ def _check_setup() -> None: try: # covalent is needed because the @electron function # executes inside `wrapper_fn` to apply deps - import covalent + _import_covalent() except ImportError as _exception: msg = "The covalent SDK is not installed in the Slurm job environment."