diff options
| -rw-r--r-- | tests/integration_tests/modules/test_ssh_import_id.py | 6 | ||||
| -rw-r--r-- | tests/integration_tests/util.py | 30 |
2 files changed, 36 insertions, 0 deletions
diff --git a/tests/integration_tests/modules/test_ssh_import_id.py b/tests/integration_tests/modules/test_ssh_import_id.py index 3db573b5..b90fe95f 100644 --- a/tests/integration_tests/modules/test_ssh_import_id.py +++ b/tests/integration_tests/modules/test_ssh_import_id.py @@ -12,6 +12,7 @@ TODO: import pytest +from tests.integration_tests.util import retry USER_DATA = """\ #cloud-config @@ -26,6 +27,11 @@ ssh_import_id: 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") 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 |