From fbf45292abb55328e36f58c9a8c2044033d939c7 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Tue, 17 Aug 2021 08:33:31 -0500
Subject: [PATCH] Add the concept of atexit hooks to Zaza

These hooks will be run after a test run completes via the
test portion of Zaza's lifecycle. An example situation where
atexit hooks can be beneficial to Zaza would be the case where
resources are created and verified in one test, and then the
thing under test is modified, and then the resources should
be re-verified. A specific example of the above would be for
a charm or payload upgrade.

An minimal example test using `atexit`:

    import zaza.atexit

    def cleanup_vm():
        for vm in cls.nova_client.servers.list():
            if vm.name == 'ins-1':
                vm.delete()
                openstack_utils.resource_removed(
                    cls.nova_client.servers,
                    vm.id,
                    msg="Waiting for the Nova VM {} to be deleted".format(vm.name))

    class VMTests(test_utils.OpenStackBaseTest):
        """Encapsulate VM tests."""

        @classmethod
        def setUpClass(cls):
            """Run class setup for running tests."""
            super(VMTests, cls).setUpClass()
            zaza.atexit(cleanup_vm)

        def test_share(self):
            instance = None
            for vm in cls.nova_client.servers.list():
                if vm.name == 'ins-1':
                    instance = vm
           id instance is None:
                instance = self.launch_guest(
                    guest_name='ins-1')
           fip = neutron_tests.floating_ips_from_instance(instance)[0]
           openstack_utils.ping_response(fip)

In the above example, a VM is spawned on the first pass through
the test case but not cleaned up. The second time through the
test case, the previous VM is used, and the verification is that
the instance is pinged.

When Zaza finishes the entire test run, it will call the atexit
function that was registered and cleanup the server.
---
 zaza/__init__.py             | 28 ++++++++++++++++++++++++++++
 zaza/charm_lifecycle/test.py |  2 ++
 2 files changed, 30 insertions(+)

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