diff --git a/zaza/__init__.py b/zaza/__init__.py index 89df308a2..8be2ce492 100644 --- a/zaza/__init__.py +++ b/zaza/__init__.py @@ -20,6 +20,7 @@ __path__ = extend_path(__path__, __name__) +_ATEXIT = [] def run(*steps): @@ -85,3 +86,30 @@ async def _run_it(): return await f(*args, **kwargs) return run(_run_it()) return _wrapper + + +def atexit(func): + """Queue the passed callable for the very end of execution. + + :param func: A callable that will be run at the end of a Zaza run + :type func: Callable + :returns: None + :raises: RuntimeError if the argument is not callable + """ + global _ATEXIT + if not callable(func): + raise RuntimeError("_atexit must be passed a Callable") + _ATEXIT.append(func) + + +def run_atexit_hooks(): + """Run the queued atexit callables. + + :returns: None + """ + global _ATEXIT + for func in _ATEXIT: + try: + func() + except Exception as e: + logging.error("Error in cleanup: {}".format(e)) diff --git a/zaza/charm_lifecycle/test.py b/zaza/charm_lifecycle/test.py index 7a6d0a320..938feb2d7 100644 --- a/zaza/charm_lifecycle/test.py +++ b/zaza/charm_lifecycle/test.py @@ -19,6 +19,7 @@ import unittest import sys +import zaza import zaza.model import zaza.global_options as global_options import zaza.charm_lifecycle.utils as utils @@ -121,6 +122,7 @@ def test(model_name, tests, test_directory=None): utils.set_base_test_dir(test_dir=test_directory) zaza.model.set_juju_model(model_name) run_test_list(tests) + zaza.run_atexit_hooks() def parse_args(args):