Ability to hot-attach NICs to preprovisioned VMs before reprovisioning (#613)

Adds the ability to run the Azure preprovisioned VMs as NIC-less and
then hot-attach them when assigned for reprovision.

The NIC on the preprovisioned VM is hot-detached as soon as it reports
ready and goes into wait for one or more interfaces to be hot-attached.
Once they are attached, cloud-init gets the expected number of NICs (in
case there are more than one) that will be attached from IMDS and waits
until all of them are attached. After all the NICs are attached,
reprovision proceeds as usual.

1 files changed, 422 insertions, 14 deletions

diff --git a/tests/unittests/test_datasource/test_azure.py b/tests/unittests/test_datasource/test_azure.py
index 534314aa..e363c1f9 100644
--- a/tests/unittests/test_datasource/test_azure.py
+++ b/tests/unittests/test_datasource/test_azure.py
@@ -11,6 +11,7 @@ from cloudinit.version import version_string as vs
 from cloudinit.tests.helpers import (
     HttprettyTestCase, CiTestCase, populate_dir, mock, wrap_and_call,
     ExitStack, resourceLocation)
+from cloudinit.sources.helpers import netlink
 
 import copy
 import crypt
@@ -78,6 +79,8 @@ def construct_valid_ovf_env(data=None, pubkeys=None,
     if platform_settings:
         for k, v in platform_settings.items():
             content += "<%s>%s</%s>\n" % (k, v, k)
+        if "PreprovisionedVMType" not in platform_settings:
+            content += """<PreprovisionedVMType i:nil="true" />"""
     content += """</PlatformSettings></wa:PlatformSettingsSection>
 </Environment>"""
 
@@ -156,6 +159,31 @@ SECONDARY_INTERFACE = {
     }
 }
 
+IMDS_NETWORK_METADATA = {
+    "interface": [
+        {
+            "macAddress": "000D3A047598",
+            "ipv6": {
+                "ipAddress": []
+            },
+            "ipv4": {
+                "subnet": [
+                    {
+                        "prefix": "24",
+                        "address": "10.0.0.0"
+                    }
+                ],
+                "ipAddress": [
+                    {
+                        "privateIpAddress": "10.0.0.4",
+                        "publicIpAddress": "104.46.124.81"
+                    }
+                ]
+            }
+        }
+    ]
+}
+
 MOCKPATH = 'cloudinit.sources.DataSourceAzure.'
 
 
@@ -366,8 +394,8 @@ class TestGetMetadataFromIMDS(HttprettyTestCase):
         self.network_md_url = dsaz.IMDS_URL + "instance?api-version=2019-06-01"
 
     @mock.patch(MOCKPATH + 'readurl')
-    @mock.patch(MOCKPATH + 'EphemeralDHCPv4')
-    @mock.patch(MOCKPATH + 'net.is_up')
+    @mock.patch(MOCKPATH + 'EphemeralDHCPv4', autospec=True)
+    @mock.patch(MOCKPATH + 'net.is_up', autospec=True)
     def test_get_metadata_does_not_dhcp_if_network_is_up(
             self, m_net_is_up, m_dhcp, m_readurl):
         """Do not perform DHCP setup when nic is already up."""
@@ -384,9 +412,66 @@ class TestGetMetadataFromIMDS(HttprettyTestCase):
             "Crawl of Azure Instance Metadata Service (IMDS) took",  # log_time
             self.logs.getvalue())
 
-    @mock.patch(MOCKPATH + 'readurl')
-    @mock.patch(MOCKPATH + 'EphemeralDHCPv4WithReporting')
+    @mock.patch(MOCKPATH + 'readurl', autospec=True)
+    @mock.patch(MOCKPATH + 'EphemeralDHCPv4')
+    @mock.patch(MOCKPATH + 'net.is_up')
+    def test_get_compute_metadata_uses_compute_url(
+            self, m_net_is_up, m_dhcp, m_readurl):
+        """Make sure readurl is called with the correct url when accessing
+        network metadata"""
+        m_net_is_up.return_value = True
+        m_readurl.return_value = url_helper.StringResponse(
+            json.dumps(IMDS_NETWORK_METADATA).encode('utf-8'))
+
+        dsaz.get_metadata_from_imds(
+            'eth0', retries=3, md_type=dsaz.metadata_type.compute)
+        m_readurl.assert_called_with(
+            "http://169.254.169.254/metadata/instance?api-version="
+            "2019-06-01", exception_cb=mock.ANY,
+            headers=mock.ANY, retries=mock.ANY,
+            timeout=mock.ANY)
+
+    @mock.patch(MOCKPATH + 'readurl', autospec=True)
+    @mock.patch(MOCKPATH + 'EphemeralDHCPv4')
+    @mock.patch(MOCKPATH + 'net.is_up')
+    def test_get_network_metadata_uses_network_url(
+            self, m_net_is_up, m_dhcp, m_readurl):
+        """Make sure readurl is called with the correct url when accessing
+        network metadata"""
+        m_net_is_up.return_value = True
+        m_readurl.return_value = url_helper.StringResponse(
+            json.dumps(IMDS_NETWORK_METADATA).encode('utf-8'))
+
+        dsaz.get_metadata_from_imds(
+            'eth0', retries=3, md_type=dsaz.metadata_type.network)
+        m_readurl.assert_called_with(
+            "http://169.254.169.254/metadata/instance/network?api-version="
+            "2019-06-01", exception_cb=mock.ANY,
+            headers=mock.ANY, retries=mock.ANY,
+            timeout=mock.ANY)
+
+    @mock.patch(MOCKPATH + 'readurl', autospec=True)
+    @mock.patch(MOCKPATH + 'EphemeralDHCPv4')
     @mock.patch(MOCKPATH + 'net.is_up')
+    def test_get_default_metadata_uses_compute_url(
+            self, m_net_is_up, m_dhcp, m_readurl):
+        """Make sure readurl is called with the correct url when accessing
+        network metadata"""
+        m_net_is_up.return_value = True
+        m_readurl.return_value = url_helper.StringResponse(
+            json.dumps(IMDS_NETWORK_METADATA).encode('utf-8'))
+
+        dsaz.get_metadata_from_imds(
+            'eth0', retries=3)
+        m_readurl.assert_called_with(
+            "http://169.254.169.254/metadata/instance?api-version="
+            "2019-06-01", exception_cb=mock.ANY,
+            headers=mock.ANY, retries=mock.ANY,
+            timeout=mock.ANY)
+
+    @mock.patch(MOCKPATH + 'readurl', autospec=True)
+    @mock.patch(MOCKPATH + 'EphemeralDHCPv4WithReporting', autospec=True)
+    @mock.patch(MOCKPATH + 'net.is_up', autospec=True)
     def test_get_metadata_performs_dhcp_when_network_is_down(
             self, m_net_is_up, m_dhcp, m_readurl):
         """Perform DHCP setup when nic is not up."""
@@ -410,7 +495,7 @@ class TestGetMetadataFromIMDS(HttprettyTestCase):
             timeout=dsaz.IMDS_TIMEOUT_IN_SECONDS)
 
     @mock.patch('cloudinit.url_helper.time.sleep')
-    @mock.patch(MOCKPATH + 'net.is_up')
+    @mock.patch(MOCKPATH + 'net.is_up', autospec=True)
     def test_get_metadata_from_imds_empty_when_no_imds_present(
             self, m_net_is_up, m_sleep):
         """Return empty dict when IMDS network metadata is absent."""
@@ -431,7 +516,7 @@ class TestGetMetadataFromIMDS(HttprettyTestCase):
 
     @mock.patch('requests.Session.request')
     @mock.patch('cloudinit.url_helper.time.sleep')
-    @mock.patch(MOCKPATH + 'net.is_up')
+    @mock.patch(MOCKPATH + 'net.is_up', autospec=True)
     def test_get_metadata_from_imds_retries_on_timeout(
             self, m_net_is_up, m_sleep, m_request):
         """Retry IMDS network metadata on timeout errors."""
@@ -801,6 +886,7 @@ scbus-1 on xpt0 bus 0
                 'sys_cfg': {}}
         dsrc = self._get_ds(data)
         expected_cfg = {
+            'PreprovisionedVMType': None,
             'PreprovisionedVm': False,
             'datasource': {'Azure': {'agent_command': 'my_command'}},
             'system_info': {'default_user': {'name': u'myuser'}}}
@@ -864,6 +950,66 @@ scbus-1 on xpt0 bus 0
         dsrc.crawl_metadata()
         self.assertEqual(1, m_report_ready.call_count)
 
+    @mock.patch(
+        'cloudinit.sources.DataSourceAzure.EphemeralDHCPv4WithReporting')
+    @mock.patch('cloudinit.sources.DataSourceAzure.util.write_file')
+    @mock.patch(
+        'cloudinit.sources.DataSourceAzure.DataSourceAzure._report_ready')
+    @mock.patch('cloudinit.sources.DataSourceAzure.DataSourceAzure._poll_imds')
+    @mock.patch(
+        'cloudinit.sources.DataSourceAzure.DataSourceAzure.'
+        '_wait_for_all_nics_ready')
+    def test_crawl_metadata_waits_for_nic_on_savable_vms(
+        self, detect_nics, poll_imds_func, report_ready_func, m_write, m_dhcp
+    ):
+        """If reprovisioning, report ready at the end"""
+        ovfenv = construct_valid_ovf_env(
+            platform_settings={"PreprovisionedVMType": "Savable",
+                               "PreprovisionedVm": "True"}
+        )
+
+        data = {
+            'ovfcontent': ovfenv,
+            'sys_cfg': {}
+        }
+        dsrc = self._get_ds(data)
+        poll_imds_func.return_value = ovfenv
+        dsrc.crawl_metadata()
+        self.assertEqual(1, report_ready_func.call_count)
+        self.assertEqual(1, detect_nics.call_count)
+
+    @mock.patch(
+        'cloudinit.sources.DataSourceAzure.EphemeralDHCPv4WithReporting')
+    @mock.patch('cloudinit.sources.DataSourceAzure.util.write_file')
+    @mock.patch(
+        'cloudinit.sources.DataSourceAzure.DataSourceAzure._report_ready')
+    @mock.patch('cloudinit.sources.DataSourceAzure.DataSourceAzure._poll_imds')
+    @mock.patch(
+        'cloudinit.sources.DataSourceAzure.DataSourceAzure.'
+        '_wait_for_all_nics_ready')
+    @mock.patch('os.path.isfile')
+    def test_detect_nics_when_marker_present(
+        self, is_file, detect_nics, poll_imds_func, report_ready_func, m_write,
+            m_dhcp):
+        """If reprovisioning, wait for nic attach if marker present"""
+
+        def is_file_ret(key):
+            return key == dsaz.REPROVISION_NIC_ATTACH_MARKER_FILE
+
+        is_file.side_effect = is_file_ret
+        ovfenv = construct_valid_ovf_env()
+
+        data = {
+            'ovfcontent': ovfenv,
+            'sys_cfg': {}
+        }
+
+        dsrc = self._get_ds(data)
+        poll_imds_func.return_value = ovfenv
+        dsrc.crawl_metadata()
+        self.assertEqual(1, report_ready_func.call_count)
+        self.assertEqual(1, detect_nics.call_count)
+
     @mock.patch('cloudinit.sources.DataSourceAzure.util.write_file')
     @mock.patch('cloudinit.sources.helpers.netlink.'
                 'wait_for_media_disconnect_connect')
@@ -1526,7 +1672,7 @@ scbus-1 on xpt0 bus 0
     @mock.patch('cloudinit.net.get_interface_mac')
     @mock.patch('cloudinit.net.get_devicelist')
     @mock.patch('cloudinit.net.device_driver')
-    @mock.patch('cloudinit.net.generate_fallback_config')
+    @mock.patch('cloudinit.net.generate_fallback_config', autospec=True)
     def test_fallback_network_config(self, mock_fallback, mock_dd,
                                      mock_devlist, mock_get_mac):
         """On absent IMDS network data, generate network fallback config."""
@@ -1561,7 +1707,7 @@ scbus-1 on xpt0 bus 0
             blacklist_drivers=['mlx4_core', 'mlx5_core'],
             config_driver=True)
 
-    @mock.patch(MOCKPATH + 'net.get_interfaces')
+    @mock.patch(MOCKPATH + 'net.get_interfaces', autospec=True)
     @mock.patch(MOCKPATH + 'util.is_FreeBSD')
     def test_blacklist_through_distro(
             self, m_is_freebsd, m_net_get_interfaces):
@@ -1583,17 +1729,17 @@ scbus-1 on xpt0 bus 0
         m_net_get_interfaces.assert_called_with(
             blacklist_drivers=dsaz.BLACKLIST_DRIVERS)
 
-    @mock.patch(MOCKPATH + 'subp.subp')
+    @mock.patch(MOCKPATH + 'subp.subp', autospec=True)
     def test_get_hostname_with_no_args(self, m_subp):
         dsaz.get_hostname()
         m_subp.assert_called_once_with(("hostname",), capture=True)
 
-    @mock.patch(MOCKPATH + 'subp.subp')
+    @mock.patch(MOCKPATH + 'subp.subp', autospec=True)
     def test_get_hostname_with_string_arg(self, m_subp):
         dsaz.get_hostname(hostname_command="hostname")
         m_subp.assert_called_once_with(("hostname",), capture=True)
 
-    @mock.patch(MOCKPATH + 'subp.subp')
+    @mock.patch(MOCKPATH + 'subp.subp', autospec=True)
     def test_get_hostname_with_iterable_arg(self, m_subp):
         dsaz.get_hostname(hostname_command=("hostname",))
         m_subp.assert_called_once_with(("hostname",), capture=True)
@@ -2224,6 +2370,29 @@ class TestPreprovisioningReadAzureOvfFlag(CiTestCase):
         ret = dsaz.read_azure_ovf(content)
         cfg = ret[2]
         self.assertFalse(cfg['PreprovisionedVm'])
+        self.assertEqual(None, cfg["PreprovisionedVMType"])
+
+    def test_read_azure_ovf_with_running_type(self):
+        """The read_azure_ovf method should set PreprovisionedVMType
+           cfg flag to Running."""
+        content = construct_valid_ovf_env(
+            platform_settings={"PreprovisionedVMType": "Running",
+                               "PreprovisionedVm": "True"})
+        ret = dsaz.read_azure_ovf(content)
+        cfg = ret[2]
+        self.assertTrue(cfg['PreprovisionedVm'])
+        self.assertEqual("Running", cfg['PreprovisionedVMType'])
+
+    def test_read_azure_ovf_with_savable_type(self):
+        """The read_azure_ovf method should set PreprovisionedVMType
+           cfg flag to Savable."""
+        content = construct_valid_ovf_env(
+            platform_settings={"PreprovisionedVMType": "Savable",
+                               "PreprovisionedVm": "True"})
+        ret = dsaz.read_azure_ovf(content)
+        cfg = ret[2]
+        self.assertTrue(cfg['PreprovisionedVm'])
+        self.assertEqual("Savable", cfg['PreprovisionedVMType'])
 
 
 @mock.patch('os.path.isfile')
@@ -2273,6 +2442,227 @@ class TestPreprovisioningShouldReprovision(CiTestCase):
         _poll_imds.assert_called_with()
 
 
+class TestPreprovisioningHotAttachNics(CiTestCase):
+
+    def setUp(self):
+        super(TestPreprovisioningHotAttachNics, self).setUp()
+        self.tmp = self.tmp_dir()
+        self.waagent_d = self.tmp_path('/var/lib/waagent', self.tmp)
+        self.paths = helpers.Paths({'cloud_dir': self.tmp})
+        dsaz.BUILTIN_DS_CONFIG['data_dir'] = self.waagent_d
+        self.paths = helpers.Paths({'cloud_dir': self.tmp})
+
+    @mock.patch('cloudinit.sources.helpers.netlink.wait_for_nic_detach_event',
+                autospec=True)
+    @mock.patch(MOCKPATH + 'util.write_file', autospec=True)
+    def test_nic_detach_writes_marker(self, m_writefile, m_detach):
+        """When we detect that a nic gets detached, we write a marker for it"""
+        dsa = dsaz.DataSourceAzure({}, distro=None, paths=self.paths)
+        nl_sock = mock.MagicMock()
+        dsa._wait_for_nic_detach(nl_sock)
+        m_detach.assert_called_with(nl_sock)
+        self.assertEqual(1, m_detach.call_count)
+        m_writefile.assert_called_with(
+            dsaz.REPROVISION_NIC_DETACHED_MARKER_FILE, mock.ANY)
+
+    @mock.patch(MOCKPATH + 'util.write_file', autospec=True)
+    @mock.patch(MOCKPATH + 'DataSourceAzure.fallback_interface')
+    @mock.patch(MOCKPATH + 'EphemeralDHCPv4WithReporting')
+    @mock.patch(MOCKPATH + 'DataSourceAzure._report_ready')
+    @mock.patch(MOCKPATH + 'DataSourceAzure._wait_for_nic_detach')
+    def test_detect_nic_attach_reports_ready_and_waits_for_detach(
+            self, m_detach, m_report_ready, m_dhcp, m_fallback_if,
+            m_writefile):
+        """Report ready first and then wait for nic detach"""
+        dsa = dsaz.DataSourceAzure({}, distro=None, paths=self.paths)
+        dsa._wait_for_all_nics_ready()
+        m_fallback_if.return_value = "Dummy interface"
+        self.assertEqual(1, m_report_ready.call_count)
+        self.assertEqual(1, m_detach.call_count)
+        self.assertEqual(1, m_writefile.call_count)
+        self.assertEqual(1, m_dhcp.call_count)
+        m_writefile.assert_called_with(dsaz.REPORTED_READY_MARKER_FILE,
+                                       mock.ANY)
+
+    @mock.patch('os.path.isfile')
+    @mock.patch(MOCKPATH + 'DataSourceAzure.fallback_interface')
+    @mock.patch(MOCKPATH + 'EphemeralDHCPv4WithReporting')
+    @mock.patch(MOCKPATH + 'DataSourceAzure._report_ready')
+    @mock.patch(MOCKPATH + 'DataSourceAzure._wait_for_nic_detach')
+    def test_detect_nic_attach_skips_report_ready_when_marker_present(
+            self, m_detach, m_report_ready, m_dhcp, m_fallback_if, m_isfile):
+        """Skip reporting ready if we already have a marker file."""
+        dsa = dsaz.DataSourceAzure({}, distro=None, paths=self.paths)
+
+        def isfile(key):
+            return key == dsaz.REPORTED_READY_MARKER_FILE
+
+        m_isfile.side_effect = isfile
+        dsa._wait_for_all_nics_ready()
+        m_fallback_if.return_value = "Dummy interface"
+        self.assertEqual(0, m_report_ready.call_count)
+        self.assertEqual(0, m_dhcp.call_count)
+        self.assertEqual(1, m_detach.call_count)
+
+    @mock.patch('os.path.isfile')
+    @mock.patch(MOCKPATH + 'DataSourceAzure.fallback_interface')
+    @mock.patch(MOCKPATH + 'EphemeralDHCPv4WithReporting')
+    @mock.patch(MOCKPATH + 'DataSourceAzure._report_ready')
+    @mock.patch(MOCKPATH + 'DataSourceAzure._wait_for_nic_detach')
+    def test_detect_nic_attach_skips_nic_detach_when_marker_present(
+            self, m_detach, m_report_ready, m_dhcp, m_fallback_if, m_isfile):
+        """Skip wait for nic detach if it already happened."""
+        dsa = dsaz.DataSourceAzure({}, distro=None, paths=self.paths)
+
+        m_isfile.return_value = True
+        dsa._wait_for_all_nics_ready()
+        m_fallback_if.return_value = "Dummy interface"
+        self.assertEqual(0, m_report_ready.call_count)
+        self.assertEqual(0, m_dhcp.call_count)
+        self.assertEqual(0, m_detach.call_count)
+
+    @mock.patch(MOCKPATH + 'DataSourceAzure.wait_for_link_up', autospec=True)
+    @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 + 'EphemeralDHCPv4')
+    @mock.patch(MOCKPATH + 'DataSourceAzure._wait_for_nic_detach')
+    @mock.patch('os.path.isfile')
+    def test_wait_for_nic_attach_if_no_fallback_interface(
+            self, m_isfile, m_detach, m_dhcpv4, m_imds, m_fallback_if,
+            m_attach, m_link_up):
+        """Wait for nic attach if we do not have a fallback interface"""
+        dsa = dsaz.DataSourceAzure({}, distro=None, paths=self.paths)
+        lease = {
+            'interface': 'eth9', 'fixed-address': '192.168.2.9',
+            'routers': '192.168.2.1', 'subnet-mask': '255.255.255.0',
+            'unknown-245': '624c3620'}
+
+        m_isfile.return_value = True
+        m_attach.return_value = "eth0"
+        dhcp_ctx = mock.MagicMock(lease=lease)
+        dhcp_ctx.obtain_lease.return_value = lease
+        m_dhcpv4.return_value = dhcp_ctx
+        m_imds.return_value = IMDS_NETWORK_METADATA
+        m_fallback_if.return_value = None
+
+        dsa._wait_for_all_nics_ready()
+
+        self.assertEqual(0, m_detach.call_count)
+        self.assertEqual(1, m_attach.call_count)
+        self.assertEqual(1, m_dhcpv4.call_count)
+        self.assertEqual(1, m_imds.call_count)
+        self.assertEqual(1, m_link_up.call_count)
+        m_link_up.assert_called_with(mock.ANY, "eth0")
+
+    @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 + 'EphemeralDHCPv4')
+    @mock.patch(MOCKPATH + 'DataSourceAzure._wait_for_nic_detach')
+    @mock.patch('os.path.isfile')
+    def test_wait_for_nic_attach_multinic_attach(
+            self, m_isfile, m_detach, m_dhcpv4, m_imds, m_fallback_if,
+            m_attach, m_link_up):
+        """Wait for nic attach if we do not have a fallback interface"""
+        dsa = dsaz.DataSourceAzure({}, distro=None, paths=self.paths)
+        lease = {
+            'interface': 'eth9', 'fixed-address': '192.168.2.9',
+            'routers': '192.168.2.1', 'subnet-mask': '255.255.255.0',
+            'unknown-245': '624c3620'}
+        m_attach_call_count = 0
+
+        def nic_attach_ret(nl_sock, nics_found):
+            nonlocal m_attach_call_count
+            if m_attach_call_count == 0:
+                m_attach_call_count = m_attach_call_count + 1
+                return "eth0"
+            return "eth1"
+
+        def network_metadata_ret(ifname, retries, type):
+            # Simulate two NICs by adding the same one twice.
+            md = IMDS_NETWORK_METADATA
+            md['interface'].append(md['interface'][0])
+            if ifname == "eth0":
+                return md
+            raise requests.Timeout('Fake connection timeout')
+
+        m_isfile.return_value = True
+        m_attach.side_effect = nic_attach_ret
+        dhcp_ctx = mock.MagicMock(lease=lease)
+        dhcp_ctx.obtain_lease.return_value = lease
+        m_dhcpv4.return_value = dhcp_ctx
+        m_imds.side_effect = network_metadata_ret
+        m_fallback_if.return_value = None
+
+        dsa._wait_for_all_nics_ready()
+
+        self.assertEqual(0, m_detach.call_count)
+        self.assertEqual(2, m_attach.call_count)
+        # DHCP and network metadata calls will only happen on the primary NIC.
+        self.assertEqual(1, m_dhcpv4.call_count)
+        self.assertEqual(1, m_imds.call_count)
+        self.assertEqual(2, m_link_up.call_count)
+
+    @mock.patch('cloudinit.distros.networking.LinuxNetworking.try_set_link_up')
+    def test_wait_for_link_up_returns_if_already_up(
+            self, m_is_link_up):
+        """Waiting for link to be up should return immediately if the link is
+           already up."""
+
+        distro_cls = distros.fetch('ubuntu')
+        distro = distro_cls('ubuntu', {}, self.paths)
+        dsa = dsaz.DataSourceAzure({}, distro=distro, paths=self.paths)
+        m_is_link_up.return_value = True
+
+        dsa.wait_for_link_up("eth0")
+        self.assertEqual(1, m_is_link_up.call_count)
+
+    @mock.patch(MOCKPATH + 'util.write_file')
+    @mock.patch('cloudinit.net.read_sys_net')
+    @mock.patch('cloudinit.distros.networking.LinuxNetworking.try_set_link_up')
+    def test_wait_for_link_up_writes_to_device_file(
+            self, m_is_link_up, m_read_sys_net, m_writefile):
+        """Waiting for link to be up should return immediately if the link is
+           already up."""
+
+        distro_cls = distros.fetch('ubuntu')
+        distro = distro_cls('ubuntu', {}, self.paths)
+        dsa = dsaz.DataSourceAzure({}, distro=distro, paths=self.paths)
+
+        callcount = 0
+
+        def linkup(key):
+            nonlocal callcount
+            if callcount == 0:
+                callcount += 1
+                return False
+            return True
+
+        m_is_link_up.side_effect = linkup
+
+        dsa.wait_for_link_up("eth0")
+        self.assertEqual(2, m_is_link_up.call_count)
+        self.assertEqual(1, m_read_sys_net.call_count)
+        self.assertEqual(2, m_writefile.call_count)
+
+    @mock.patch('cloudinit.sources.helpers.netlink.'
+                'create_bound_netlink_socket')
+    def test_wait_for_all_nics_ready_raises_if_socket_fails(self, m_socket):
+        """Waiting for all nics should raise exception if netlink socket
+           creation fails."""
+
+        m_socket.side_effect = netlink.NetlinkCreateSocketError
+        distro_cls = distros.fetch('ubuntu')
+        distro = distro_cls('ubuntu', {}, self.paths)
+        dsa = dsaz.DataSourceAzure({}, distro=distro, paths=self.paths)
+
+        self.assertRaises(netlink.NetlinkCreateSocketError,
+                          dsa._wait_for_all_nics_ready)
+        # dsa._wait_for_all_nics_ready()
+
+
 @mock.patch('cloudinit.net.dhcp.EphemeralIPv4Network')
 @mock.patch('cloudinit.net.dhcp.maybe_perform_dhcp_discovery')
 @mock.patch('cloudinit.sources.helpers.netlink.'
@@ -2330,6 +2720,24 @@ class TestPreprovisioningPollIMDS(CiTestCase):
         self.assertEqual(3, m_dhcpv4.call_count, 'Expected 3 DHCP calls')
         self.assertEqual(4, self.tries, 'Expected 4 total reads from IMDS')
 
+    @mock.patch('os.path.isfile')
+    def test_poll_imds_skips_dhcp_if_ctx_present(
+            self, m_isfile, report_ready_func, fake_resp, m_media_switch,
+            m_dhcp, m_net):
+        """The poll_imds function should reuse the dhcp ctx if it is already
+           present. This happens when we wait for nic to be hot-attached before
+           polling for reprovisiondata. Note that if this ctx is set when
+           _poll_imds is called, then it is not expected to be waiting for
+           media_disconnect_connect either."""
+        report_file = self.tmp_path('report_marker', self.tmp)
+        m_isfile.return_value = True
+        dsa = dsaz.DataSourceAzure({}, distro=None, paths=self.paths)
+        dsa._ephemeral_dhcp_ctx = "Dummy dhcp ctx"
+        with mock.patch(MOCKPATH + 'REPORTED_READY_MARKER_FILE', report_file):
+            dsa._poll_imds()
+        self.assertEqual(0, m_dhcp.call_count)
+        self.assertEqual(0, m_media_switch.call_count)
+
     def test_does_not_poll_imds_report_ready_when_marker_file_exists(
             self, m_report_ready, m_request, m_media_switch, m_dhcp, m_net):
         """poll_imds should not call report ready when the reported ready
@@ -2390,7 +2798,7 @@ class TestPreprovisioningPollIMDS(CiTestCase):
 @mock.patch(MOCKPATH + 'util.is_FreeBSD')
 @mock.patch('cloudinit.sources.helpers.netlink.'
             'wait_for_media_disconnect_connect')
-@mock.patch('cloudinit.net.dhcp.EphemeralIPv4Network')
+@mock.patch('cloudinit.net.dhcp.EphemeralIPv4Network', autospec=True)
 @mock.patch('cloudinit.net.dhcp.maybe_perform_dhcp_discovery')
 @mock.patch('requests.Session.request')
 class TestAzureDataSourcePreprovisioning(CiTestCase):
@@ -2412,7 +2820,7 @@ class TestAzureDataSourcePreprovisioning(CiTestCase):
         m_dhcp.return_value = [{
             'interface': 'eth9', 'fixed-address': '192.168.2.9',
             'routers': '192.168.2.1', 'subnet-mask': '255.255.255.0'}]
-        url = 'http://{0}/metadata/reprovisiondata?api-version=2017-04-02'
+        url = 'http://{0}/metadata/reprovisiondata?api-version=2019-06-01'
         host = "169.254.169.254"
         full_url = url.format(host)
         m_request.return_value = mock.MagicMock(status_code=200, text="ovf",
@@ -2445,7 +2853,7 @@ class TestAzureDataSourcePreprovisioning(CiTestCase):
             'interface': 'eth9', 'fixed-address': '192.168.2.9',
             'routers': '192.168.2.1', 'subnet-mask': '255.255.255.0',
             'unknown-245': '624c3620'}]
-        url = 'http://{0}/metadata/reprovisiondata?api-version=2017-04-02'
+        url = 'http://{0}/metadata/reprovisiondata?api-version=2019-06-01'
         host = "169.254.169.254"
         full_url = url.format(host)
         hostname = "myhost"