From 68eb4f482a84443412d4b6d6b68e2d1339520288 Mon Sep 17 00:00:00 2001
From: Frode Nordahl <frode.nordahl@canonical.com>
Date: Wed, 13 Apr 2022 15:36:03 +0200
Subject: [PATCH] Update zaza/__init__.py from zaza (#750)

The zaza and zaza-openstack-tests project are currently installed
within the same package.  A side effect of that is that the
zaza/__init__.py from zaza is overwritten by the one in
zaza-openstack-tests.

We ought to find a better way of doing this but we're currently
blocked on using zaza + zaza-openstack-tests on systems with
Python 3.10 installed, and updating this file unblocks us.

Update zaza/__init__.py from openstack-charmers/zaza@c48c955ef7dd8393b54ac96bf72689f757df1364
---
 zaza/__init__.py | 66 +++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 63 insertions(+), 3 deletions(-)

diff --git a/zaza/__init__.py b/zaza/__init__.py
index e10a793..7f215a4 100644
--- a/zaza/__init__.py
+++ b/zaza/__init__.py
@@ -12,9 +12,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# __NOTE__
+#
+# Whenever this file is changed, make sure to update the copy of it in
+# ``zaza-openstack-tests``.
+#
+# The ``zaza`` and ``zaza-openstack-tests`` projects are related, and currently
+# the latter is installed as a package inside the former.  As a consequence
+# ``zaza-openstack-tests`` needs to carry a copy of this file
+# (``zaza/__init__.py``) as this file will be overwritten by the copy in
+# ``zaza-openstack-tests`` on install.
+#
+# We of course want a better solution to this, but in the interest of time
+# this note is left here until we get around to fixing it properly.
+#
+# __NOTE__
+
 """Functions to support converting async function to a sync equivalent."""
 import asyncio
+import logging
 from pkgutil import extend_path
+from sys import version_info
 
 
 __path__ = extend_path(__path__, __name__)
@@ -23,22 +41,64 @@ __path__ = extend_path(__path__, __name__)
 def run(*steps):
     """Run the given steps in an asyncio loop.
 
-    :returns: The result of the asyncio.Task
+    If the tasks spawns other future (tasks) then these are also cleaned up
+    after each step is performed.
+
+    :returns: The result of the last asyncio.Task
     :rtype: Any
     """
     if not steps:
         return
-    loop = asyncio.get_event_loop()
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = asyncio.new_event_loop()
+    except AttributeError:
+        # Remove once support for Python 3.6 is dropped
+        loop = asyncio.get_event_loop()
 
     for step in steps:
         task = loop.create_task(step)
-        loop.run_until_complete(asyncio.wait([task], loop=loop))
+        loop.run_until_complete(asyncio.wait([task]))
+
+        # Let's also cancel any remaining tasks:
+        while True:
+            # issue #445 - asyncio.Task.all_tasks() deprecated in 3.7
+            if version_info.major == 3 and version_info.minor >= 7:
+                try:
+                    tasklist = asyncio.all_tasks()
+                except RuntimeError:
+                    # no running event loop
+                    break
+            else:
+                tasklist = asyncio.Task.all_tasks()
+            pending_tasks = [p for p in tasklist if not p.done()]
+            if pending_tasks:
+                logging.info(
+                    "async -> sync. cleaning up pending tasks: len: {}"
+                    .format(len(pending_tasks)))
+                for pending_task in pending_tasks:
+                    pending_task.cancel()
+                    try:
+                        loop.run_until_complete(pending_task)
+                    except asyncio.CancelledError:
+                        pass
+                    except Exception as e:
+                        logging.error(
+                            "A pending task caused an exception: {}"
+                            .format(str(e)))
+            else:
+                break
+
     return task.result()
 
 
 def sync_wrapper(f):
     """Convert the given async function into a sync function.
 
+    This is only to be called from sync code and it runs all tasks (and cancels
+    all tasks at the end of each run) for the code that is being given.
+
     :returns: The de-async'd function
     :rtype: function
     """