From c041042fe220fa5de33f8006734c66acb8255f11 Mon Sep 17 00:00:00 2001
From: Alex Kavanagh <alex.kavanagh@canonical.com>
Date: Mon, 8 Mar 2021 10:50:41 +0000
Subject: [PATCH] Add function to just retry on keystone ConnectFailure

The main failure that seems to occur with clients is the
ConnectFailure, which according to the docs, is retry-able.  Thus
provide a function that adds that exception condition automatically.
Also fix the import problem in octavia test for the object retrier that
is being used to validate the ObjectRetrier feature.
---
 unit_tests/utilities/test_utilities.py      | 20 ++++++++++++++++++
 zaza/openstack/charm_tests/octavia/tests.py |  2 +-
 zaza/openstack/utilities/__init__.py        | 23 +++++++++++++++++++++
 3 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/unit_tests/utilities/test_utilities.py b/unit_tests/utilities/test_utilities.py
index e961de1..a78600c 100644
--- a/unit_tests/utilities/test_utilities.py
+++ b/unit_tests/utilities/test_utilities.py
@@ -164,3 +164,23 @@ class TestObjectRetrier(ut_utils.BaseTestCase):
         # that.
         mock_sleep.assert_has_calls([mock.call(5),
                                      mock.call(5)])
+
+    @mock.patch("time.sleep")
+    def test_retry_on_connect_failure(self, mock_sleep):
+
+        class A:
+
+            def func1(self):
+                raise SomeException()
+
+            def func2(self):
+                raise utilities.ConnectFailure()
+
+        a = A()
+        wrapped_a = utilities.retry_on_connect_failure(a, num_retries=2)
+        with self.assertRaises(SomeException):
+            wrapped_a.func1()
+        mock_sleep.assert_not_called()
+        with self.assertRaises(utilities.ConnectFailure):
+            wrapped_a.func2()
+        mock_sleep.assert_has_calls([mock.call(5)])
diff --git a/zaza/openstack/charm_tests/octavia/tests.py b/zaza/openstack/charm_tests/octavia/tests.py
index d6bae64..9a71b49 100644
--- a/zaza/openstack/charm_tests/octavia/tests.py
+++ b/zaza/openstack/charm_tests/octavia/tests.py
@@ -25,7 +25,7 @@ import osc_lib.exceptions
 import zaza.openstack.charm_tests.test_utils as test_utils
 import zaza.openstack.utilities.openstack as openstack_utils
 
-from zaza.openstack import ObjectRetrier
+from zaza.openstack.utilities import ObjectRetrier
 
 
 class CharmOperationTest(test_utils.OpenStackBaseTest):
diff --git a/zaza/openstack/utilities/__init__.py b/zaza/openstack/utilities/__init__.py
index 5011689..6db3ce3 100644
--- a/zaza/openstack/utilities/__init__.py
+++ b/zaza/openstack/utilities/__init__.py
@@ -17,6 +17,8 @@
 
 import time
 
+from keystoneauth1.exceptions.connection import ConnectFailure
+
 
 class ObjectRetrier(object):
     """An automatic retrier for an object.
@@ -141,3 +143,24 @@ class ObjectRetrier(object):
                 wait = wait * backoff
                 if wait > max_interval:
                     wait = max_interval
+
+
+def retry_on_connect_failure(client, **kwargs):
+    """Retry an object that eventually gets resolved to a call.
+
+    Specifically, this uses ObjectRetrier but only against the
+    keystoneauth1.exceptions.connection.ConnectFailure exeception.
+
+    :params client: the object that may throw and exception when called.
+    :type client: Any
+    :params **kwargs: the arguments supplied to the ObjectRetrier init method
+    :type **kwargs: Dict[Any]
+    :returns: client wrapped in an ObjectRetrier instance
+    :rtype: ObjectRetrier[client]
+    """
+    kwcopy = kwargs.copy()
+    if 'retry_exceptions' not in kwcopy:
+        kwcopy['retry_exceptions'] = []
+    if ConnectFailure not in kwcopy['retry_exceptions']:
+        kwcopy['retry_exceptions'].append(ConnectFailure)
+    return ObjectRetrier(client, **kwcopy)