Develop-Python
diff --git a/‎tests/integration_tests/modules/test_ssh_import_id.py‎
Lines changed: 6 additions & 0 deletions b/‎tests/integration_tests/modules/test_ssh_import_id.py‎
Lines changed: 6 additions & 0 deletions
diff --git a/‎tests/integration_tests/util.py‎
Lines changed: 30 additions & 0 deletions b/‎tests/integration_tests/util.py‎
Lines changed: 30 additions & 0 deletions
@@ -12,6 +12,7 @@
 
 import pytest
 
+from tests.integration_tests.util import retry
 
 USER_DATA = """\
 #cloud-config
@@ -26,6 +27,11 @@
 class TestSshImportId:
 
     @pytest.mark.user_data(USER_DATA)
+    # Retry is needed here because ssh import id is one of the last modules
+    # run, and it fires off a web request, then continues with the rest of
+    # cloud-init. It is possible cloud-init's status is "done" before the
+    # id's have been fully imported.
+    @retry(tries=30, delay=1)
     def test_ssh_import_id(self, client):
         ssh_output = client.read_from_file(
             "/home/ubuntu/.ssh/authorized_keys")
 
@@ -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