Fix retrying.

2 files changed, 18 insertions, 2 deletions

diff --git a/cloudinit/sources/helpers/azure.py b/cloudinit/sources/helpers/azure.py
index dfdfa7c2..2ce728f5 100644
--- a/cloudinit/sources/helpers/azure.py
+++ b/cloudinit/sources/helpers/azure.py
@@ -242,14 +242,19 @@ class WALinuxAgentShim(object):
         self.openssl_manager = OpenSSLManager()
         http_client = AzureEndpointHttpClient(self.openssl_manager.certificate)
         LOG.info('Registering with Azure...')
-        for i in range(10):
+        attempts = 0
+        while True:
             try:
                 response = http_client.get(
                     'http://{}/machine/?comp=goalstate'.format(self.endpoint))
             except Exception:
-                time.sleep(i + 1)
+                if attempts < 10:
+                    time.sleep(attempts + 1)
+                else:
+                    raise
             else:
                 break
+            attempts += 1
         LOG.debug('Successfully fetched GoalState XML.')
         goal_state = GoalState(response.contents, http_client)
         public_keys = []
diff --git a/tests/unittests/test_datasource/test_azure_helper.py b/tests/unittests/test_datasource/test_azure_helper.py
index 5fac2ade..23bc997c 100644
--- a/tests/unittests/test_datasource/test_azure_helper.py
+++ b/tests/unittests/test_datasource/test_azure_helper.py
@@ -323,6 +323,8 @@ class TestWALinuxAgentShim(TestCase):
             mock.patch.object(azure_helper, 'iid_from_shared_config_content'))
         self.OpenSSLManager = patches.enter_context(
             mock.patch.object(azure_helper, 'OpenSSLManager'))
+        patches.enter_context(
+            mock.patch.object(azure_helper.time, 'sleep', mock.MagicMock()))
 
     def test_http_client_uses_certificate(self):
         shim = azure_helper.WALinuxAgentShim()
@@ -402,6 +404,15 @@ class TestWALinuxAgentShim(TestCase):
         self.assertEqual(
             1, self.OpenSSLManager.return_value.clean_up.call_count)
 
+    def test_failure_to_fetch_goalstate_bubbles_up(self):
+        class SentinelException(Exception):
+            pass
+        self.AzureEndpointHttpClient.return_value.get.side_effect = (
+            SentinelException)
+        shim = azure_helper.WALinuxAgentShim()
+        self.assertRaises(SentinelException,
+                          shim.register_with_azure_and_fetch_data)
+
 
 class TestGetMetadataFromFabric(TestCase):