Add the concept of atexit hooks to Zaza

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))

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):