vdk-core: Test vdk_exit and vdk_exception hooks are called (#791)

Signed-off-by: David Laing <[email protected]>

projects/vdk-core/src/vdk/api/data_job.py

@@ -16,9 +16,6 @@ class IStandaloneDataJob:
 
     NB: The following plugin hooks are not triggered:
         * CoreHookSpecs.vdk_command_line
-        * JobRunHookSpecs.run_step
-        * CoreHookSpecs.vdk_exception
-        * CoreHookSpecs.vdk_exit
 
     Sample usage::
 
@@ -81,5 +78,6 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
 
         Triggers plugin lifecycle hooks:
             * finalize_job
+            * vdk_exit
         """
         pass

projects/vdk-core/src/vdk/internal/builtin_plugins/run/standalone_data_job.py

@@ -3,9 +3,11 @@
 from __future__ import annotations
 
 import pathlib
+import sys
 from pathlib import Path
 from typing import cast
 
+from click import ClickException
 from vdk.api.data_job import IStandaloneDataJob
 from vdk.api.job_input import IJobInput
 from vdk.api.plugin.core_hook_spec import CoreHookSpecs
@@ -72,59 +74,88 @@ def __init__(
             job_args = {}
         self._job_args = job_args
 
-        plugin_registry = PluginRegistry()
-        plugin_registry.add_hook_specs(InternalHookSpecs)
+        self._plugin_registry = PluginRegistry()
+        self._plugin_registry.add_hook_specs(InternalHookSpecs)
         for plugin in extra_plugins:
-            plugin_registry.load_plugin_with_hooks_impl(plugin)
-        plugin_registry.load_plugins_from_setuptools_entrypoints()
-        plugin_registry.add_hook_specs(CoreHookSpecs)
-        plugin_registry.load_plugin_with_hooks_impl(builtin_hook_impl, "core-plugin")
+            self._plugin_registry.load_plugin_with_hooks_impl(plugin)
+        self._plugin_registry.load_plugins_from_setuptools_entrypoints()
+        self._plugin_registry.add_hook_specs(CoreHookSpecs)
+        self._plugin_registry.load_plugin_with_hooks_impl(
+            builtin_hook_impl, "core-plugin"
+        )
 
-        plugin_registry.hook().vdk_start.call_historic(
-            kwargs=dict(plugin_registry=plugin_registry, command_line_args=[])
+        cast(CoreHookSpecs, self._plugin_registry.hook()).vdk_start.call_historic(
+            kwargs=dict(plugin_registry=self._plugin_registry, command_line_args=[])
         )
 
         conf_builder = ConfigurationBuilder()
-        cast(CoreHookSpecs, plugin_registry.hook()).vdk_configure(
+        cast(CoreHookSpecs, self._plugin_registry.hook()).vdk_configure(
             config_builder=conf_builder
         )
         configuration = conf_builder.build()
 
-        core_context = CoreContext(plugin_registry, configuration, StateStore())
+        core_context = CoreContext(self._plugin_registry, configuration, StateStore())
         core_context.plugin_registry.load_plugin_with_hooks_impl(
             ExecutionTrackingPlugin()
         )
-        cast(CoreHookSpecs, plugin_registry.hook()).vdk_initialize(context=core_context)
+        cast(CoreHookSpecs, self._plugin_registry.hook()).vdk_initialize(
+            context=core_context
+        )
 
-        self._plugin_hook = cast(JobRunHookSpecs, plugin_registry.hook())
         self._core_context = core_context
         self._job_context = None
 
     def __enter__(self) -> IJobInput:
-        self._core_context.plugin_registry.load_plugin_with_hooks_impl(
-            NoOpStepDataJobHookImplPlugin(), NoOpStepDataJobHookImplPlugin.__name__
-        )
+        try:
+            self._core_context.plugin_registry.load_plugin_with_hooks_impl(
+                NoOpStepDataJobHookImplPlugin(), NoOpStepDataJobHookImplPlugin.__name__
+            )
 
-        from vdk.internal.builtin_plugins.templates.template_impl import TemplatesImpl
+            from vdk.internal.builtin_plugins.templates.template_impl import (
+                TemplatesImpl,
+            )
 
-        self._job_context = JobContext(
-            name=self._name,
-            job_directory=self._data_job_directory,
-            core_context=self._core_context,
-            job_args=JobArguments(self._job_args),
-            templates=TemplatesImpl(
-                job_name=self._name, core_context=self._core_context
-            ),
-        )
+            self._job_context = JobContext(
+                name=self._name,
+                job_directory=self._data_job_directory,
+                core_context=self._core_context,
+                job_args=JobArguments(self._job_args),
+                templates=TemplatesImpl(
+                    job_name=self._name, core_context=self._core_context
+                ),
+            )
 
-        # We need to call both the initialize_job and run_job hooks to get a full initialized JobInput
-        self._plugin_hook.initialize_job(context=self._job_context)
-        self._plugin_hook.run_job(context=self._job_context)
+            # We need to call both the initialize_job and run_job hooks to get a fully initialized JobInput
+            cast(JobRunHookSpecs, self._plugin_registry.hook()).initialize_job(
+                context=self._job_context
+            )
+            cast(JobRunHookSpecs, self._plugin_registry.hook()).run_job(
+                context=self._job_context
+            )
 
-        return self._job_context.job_input
+            return self._job_context.job_input
+        except:
+            # Ensure __exit__ gets called so the vdk_exception hook can be triggered
+            self.__exit__(*sys.exc_info())
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
-        self._plugin_hook.finalize_job(context=self._job_context)
+        exit_code = 0
+        if exc_type:
+            cast(CoreHookSpecs, self._plugin_registry.hook()).vdk_exception(
+                exception=exc_value
+            )
+            exit_code = (
+                exc_value.exit_code if isinstance(exc_value, ClickException) else 1
+            )
+        else:
+            cast(JobRunHookSpecs, self._plugin_registry.hook()).finalize_job(
+                context=self._job_context
+            )
+        cast(CoreHookSpecs, self._plugin_registry.hook()).vdk_exit(
+            context=self._core_context, exit_code=exit_code
+        )
+        # Ensure the exception can be processed normally by calling code
+        return False
 
 
 class StandaloneDataJobFactory:

projects/vdk-core/tests/functional/run/test_run_standalone_data_job.py

@@ -41,22 +41,24 @@ def test_standalone_data_job_executes_hooks_in_expected_order(
 ):
     assert (
         ", ".join(executed_standalone_data_job.hook_tracker.calls)
-        == "vdk_start, vdk_configure, vdk_initialize, initialize_job, run_job, run_step, finalize_job"
+        == "vdk_start, vdk_configure, vdk_initialize, initialize_job, run_job, run_step, finalize_job, vdk_exit"
     )
 
 
-@pytest.mark.skip(
-    reason="Not sure whether vdk_exit should be called, given we're not executing from the command line"
-)
-def test_standalone_data_job_executes_vdk_exit_hooks(executed_standalone_data_job):
-    assert "vdk_exit" in executed_standalone_data_job.hook_tracker.calls
-
-
 def test_standalone_data_job_does_not_call_command_line_hook(
     executed_standalone_data_job,
 ):
     assert "vdk_command_line" not in executed_standalone_data_job.hook_tracker.calls
 
 
-def test_standalone_data_job_does_not_call_exception_hook(executed_standalone_data_job):
-    assert "vdk_exception" not in executed_standalone_data_job.hook_tracker.calls
+def test_standalone_data_job_does_calls_exception_hook_after_exception_thrown():
+    my_hook_tracker = util.HookCallTracker()
+
+    with pytest.raises(ValueError) as exc_info:
+        with StandaloneDataJobFactory.create(
+            data_job_directory=Path(__file__), extra_plugins=[my_hook_tracker]
+        ) as job_input:
+            raise ValueError("Test Error")
+
+    assert "vdk_exception" in my_hook_tracker.calls
+    assert str(exc_info.value) == "Test Error"

projects/vdk-core/tests/functional/run/util.py

@@ -24,46 +24,51 @@ def job_path(job_name: str) -> str:
 
 
 class HookCallTracker:
-    calls = []
+    @property
+    def calls(self):
+        return self._calls
+
+    def __init__(self):
+        self._calls = []
 
     @hookimpl
     def vdk_start(
         self, plugin_registry: IPluginRegistry, command_line_args: List
     ) -> None:
-        self.calls.append("vdk_start")
+        self._calls.append("vdk_start")
 
     @hookimpl
     def vdk_command_line(self, root_command: click.Group) -> None:
-        self.calls.append("vdk_command_line")
+        self._calls.append("vdk_command_line")
 
     @hookimpl
     def vdk_configure(self, config_builder: "ConfigurationBuilder") -> None:
-        self.calls.append("vdk_configure")
+        self._calls.append("vdk_configure")
 
     @hookimpl
     def vdk_initialize(self, context: CoreContext) -> None:
-        self.calls.append("vdk_initialize")
+        self._calls.append("vdk_initialize")
 
     @hookimpl
     def vdk_exception(self, exception: Exception) -> bool:
-        self.calls.append("vdk_exception")
+        self._calls.append("vdk_exception")
 
     @hookimpl
     def vdk_exit(self, context: CoreContext, exit_code: int) -> None:
-        self.calls.append("vdk_exit")
+        self._calls.append("vdk_exit")
 
     @hookimpl
     def initialize_job(self, context: JobContext) -> None:
-        self.calls.append("initialize_job")
+        self._calls.append("initialize_job")
 
     @hookimpl
     def run_job(self, context: JobContext) -> Optional[ExecutionResult]:
-        self.calls.append("run_job")
+        self._calls.append("run_job")
 
     @hookimpl
     def run_step(self, context: JobContext, step: Step) -> Optional[StepResult]:
-        self.calls.append("run_step")
+        self._calls.append("run_step")
 
     @hookimpl
     def finalize_job(self, context: JobContext) -> None:
-        self.calls.append("finalize_job")
+        self._calls.append("finalize_job")