Add flexibility to IMDS api-version (#793)

Add flexibility to IMDS api-version by having both a desired IMDS
api-version and a minimum api-version. The desired api-version will
be used first, and if that fails it will fall back to the minimum
api-version.

1 files changed, 39 insertions, 3 deletions

diff --git a/tests/unittests/test_datasource/test_azure.py b/tests/unittests/test_datasource/test_azure.py
index f597c723..dedebeb1 100644
--- a/tests/unittests/test_datasource/test_azure.py
+++ b/tests/unittests/test_datasource/test_azure.py
@@ -408,7 +408,9 @@ class TestGetMetadataFromIMDS(HttprettyTestCase):
 
     def setUp(self):
         super(TestGetMetadataFromIMDS, self).setUp()
-        self.network_md_url = dsaz.IMDS_URL + "instance?api-version=2019-06-01"
+        self.network_md_url = "{}/instance?api-version=2019-06-01".format(
+            dsaz.IMDS_URL
+        )
 
     @mock.patch(MOCKPATH + 'readurl')
     @mock.patch(MOCKPATH + 'EphemeralDHCPv4', autospec=True)
@@ -518,7 +520,7 @@ class TestGetMetadataFromIMDS(HttprettyTestCase):
         """Return empty dict when IMDS network metadata is absent."""
         httpretty.register_uri(
             httpretty.GET,
-            dsaz.IMDS_URL + 'instance?api-version=2017-12-01',
+            dsaz.IMDS_URL + '/instance?api-version=2017-12-01',
             body={}, status=404)
 
         m_net_is_up.return_value = True  # skips dhcp
@@ -1877,6 +1879,40 @@ scbus-1 on xpt0 bus 0
         ssh_keys = dsrc.get_public_ssh_keys()
         self.assertEqual(ssh_keys, ['key2'])
 
+    @mock.patch(MOCKPATH + 'get_metadata_from_imds')
+    def test_imds_api_version_wanted_nonexistent(
+            self,
+            m_get_metadata_from_imds):
+        def get_metadata_from_imds_side_eff(*args, **kwargs):
+            if kwargs['api_version'] == dsaz.IMDS_VER_WANT:
+                raise url_helper.UrlError("No IMDS version", code=400)
+            return NETWORK_METADATA
+        m_get_metadata_from_imds.side_effect = get_metadata_from_imds_side_eff
+        sys_cfg = {'datasource': {'Azure': {'apply_network_config': True}}}
+        odata = {'HostName': "myhost", 'UserName': "myuser"}
+        data = {
+            'ovfcontent': construct_valid_ovf_env(data=odata),
+            'sys_cfg': sys_cfg
+        }
+        dsrc = self._get_ds(data)
+        dsrc.get_data()
+        self.assertIsNotNone(dsrc.metadata)
+        self.assertTrue(dsrc.failed_desired_api_version)
+
+    @mock.patch(
+        MOCKPATH + 'get_metadata_from_imds', return_value=NETWORK_METADATA)
+    def test_imds_api_version_wanted_exists(self, m_get_metadata_from_imds):
+        sys_cfg = {'datasource': {'Azure': {'apply_network_config': True}}}
+        odata = {'HostName': "myhost", 'UserName': "myuser"}
+        data = {
+            'ovfcontent': construct_valid_ovf_env(data=odata),
+            'sys_cfg': sys_cfg
+        }
+        dsrc = self._get_ds(data)
+        dsrc.get_data()
+        self.assertIsNotNone(dsrc.metadata)
+        self.assertFalse(dsrc.failed_desired_api_version)
+
 
 class TestAzureBounce(CiTestCase):
 
@@ -2657,7 +2693,7 @@ class TestPreprovisioningHotAttachNics(CiTestCase):
     @mock.patch(MOCKPATH + 'DataSourceAzure.wait_for_link_up')
     @mock.patch('cloudinit.sources.helpers.netlink.wait_for_nic_attach_event')
     @mock.patch('cloudinit.sources.net.find_fallback_nic')
-    @mock.patch(MOCKPATH + 'get_metadata_from_imds')
+    @mock.patch(MOCKPATH + 'DataSourceAzure.get_imds_data_with_api_fallback')
     @mock.patch(MOCKPATH + 'EphemeralDHCPv4')
     @mock.patch(MOCKPATH + 'DataSourceAzure._wait_for_nic_detach')
     @mock.patch('os.path.isfile')