testing: fix test_ssh_import_id.py (#954)

test_ssh_import_id.py occassionally fails because cloud-init finishes
before the keys have been fully imported. A retry has been added to the
test.

1 files changed, 30 insertions, 0 deletions

diff --git a/tests/integration_tests/util.py b/tests/integration_tests/util.py
index ce62ffc8..80430eab 100644
--- a/tests/integration_tests/util.py
+++ b/tests/integration_tests/util.py
@@ -1,3 +1,4 @@
+import functools
 import logging
 import multiprocessing
 import os
@@ -64,3 +65,32 @@ def get_test_rsa_keypair(key_name: str = 'test1') -> key_pair:
     with private_key_path.open() as private_file:
         private_key = private_file.read()
     return key_pair(public_key, private_key)
+
+
+def retry(*, tries: int = 30, delay: int = 1):
+    """Decorator for retries.
+
+    Retry a function until code no longer raises an exception or
+    max tries is reached.
+
+    Example:
+      @retry(tries=5, delay=1)
+      def try_something_that_may_not_be_ready():
+          ...
+    """
+    def _retry(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            last_error = None
+            for _ in range(tries):
+                try:
+                    func(*args, **kwargs)
+                    break
+                except Exception as e:
+                    last_error = e
+                    time.sleep(delay)
+            else:
+                if last_error:
+                    raise last_error
+        return wrapper
+    return _retry