DataSourceAzure: write marker file after report ready in preprovisioning (#590)

DataSourceAzure previously writes the preprovisioning
reported ready marker file before it goes through the
report ready workflow. On certain VM instances, the
marker file is successfully written but then reporting
ready fails.

Upon rare VM reboots by the platform, cloud-init sees
that the report ready marker file already exists.
The existence of this marker file tells cloud-init
not to report ready again (because it mistakenly
assumes that it already reported ready in
preprovisioning).

In this scenario, cloud-init instead erroneously
takes the reprovisioning workflow instead of
reporting ready again.

2 files changed, 70 insertions, 19 deletions

diff --git a/tests/unittests/test_datasource/test_azure.py b/tests/unittests/test_datasource/test_azure.py
index 3b9456f7..56c1cf18 100644
--- a/tests/unittests/test_datasource/test_azure.py
+++ b/tests/unittests/test_datasource/test_azure.py
@@ -1161,6 +1161,19 @@ scbus-1 on xpt0 bus 0
         dsrc = self._get_ds({'ovfcontent': xml})
         dsrc.get_data()
 
+    def test_dsaz_report_ready_returns_true_when_report_succeeds(
+            self):
+        dsrc = self._get_ds({'ovfcontent': construct_valid_ovf_env()})
+        dsrc.ds_cfg['agent_command'] = '__builtin__'
+        self.assertTrue(dsrc._report_ready(lease=mock.MagicMock()))
+
+    def test_dsaz_report_ready_returns_false_and_does_not_propagate_exc(
+            self):
+        dsrc = self._get_ds({'ovfcontent': construct_valid_ovf_env()})
+        dsrc.ds_cfg['agent_command'] = '__builtin__'
+        self.get_metadata_from_fabric.side_effect = Exception
+        self.assertFalse(dsrc._report_ready(lease=mock.MagicMock()))
+
     def test_exception_fetching_fabric_data_doesnt_propagate(self):
         """Errors communicating with fabric should warn, but return True."""
         dsrc = self._get_ds({'ovfcontent': construct_valid_ovf_env()})
@@ -2041,7 +2054,7 @@ class TestPreprovisioningPollIMDS(CiTestCase):
     @mock.patch('time.sleep', mock.MagicMock())
     @mock.patch(MOCKPATH + 'EphemeralDHCPv4')
     def test_poll_imds_re_dhcp_on_timeout(self, m_dhcpv4, report_ready_func,
-                                          fake_resp, m_media_switch, m_dhcp,
+                                          m_request, m_media_switch, m_dhcp,
                                           m_net):
         """The poll_imds will retry DHCP on IMDS timeout."""
         report_file = self.tmp_path('report_marker', self.tmp)
@@ -2070,7 +2083,7 @@ class TestPreprovisioningPollIMDS(CiTestCase):
             # Third try should succeed and stop retries or redhcp
             return mock.MagicMock(status_code=200, text="good", content="good")
 
-        fake_resp.side_effect = fake_timeout_once
+        m_request.side_effect = fake_timeout_once
 
         dsa = dsaz.DataSourceAzure({}, distro=mock.Mock(), paths=self.paths)
         with mock.patch(MOCKPATH + 'REPORTED_READY_MARKER_FILE', report_file):
@@ -2080,11 +2093,10 @@ 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')
 
-    def test_poll_imds_report_ready_false(self,
-                                          report_ready_func, fake_resp,
-                                          m_media_switch, m_dhcp, m_net):
-        """The poll_imds should not call reporting ready
-           when flag is false"""
+    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
+        marker file exists"""
         report_file = self.tmp_path('report_marker', self.tmp)
         write_file(report_file, content='dont run report_ready :)')
         m_dhcp.return_value = [{
@@ -2095,11 +2107,49 @@ class TestPreprovisioningPollIMDS(CiTestCase):
         dsa = dsaz.DataSourceAzure({}, distro=mock.Mock(), paths=self.paths)
         with mock.patch(MOCKPATH + 'REPORTED_READY_MARKER_FILE', report_file):
             dsa._poll_imds()
-        self.assertEqual(report_ready_func.call_count, 0)
+        self.assertEqual(m_report_ready.call_count, 0)
+
+    def test_poll_imds_report_ready_success_writes_marker_file(
+            self, m_report_ready, m_request, m_media_switch, m_dhcp, m_net):
+        """poll_imds should write the report_ready marker file if
+        reporting ready succeeds"""
+        report_file = self.tmp_path('report_marker', self.tmp)
+        m_dhcp.return_value = [{
+            'interface': 'eth9', 'fixed-address': '192.168.2.9',
+            'routers': '192.168.2.1', 'subnet-mask': '255.255.255.0',
+            'unknown-245': '624c3620'}]
+        m_media_switch.return_value = None
+        dsa = dsaz.DataSourceAzure({}, distro=None, paths=self.paths)
+        self.assertFalse(os.path.exists(report_file))
+        with mock.patch(MOCKPATH + 'REPORTED_READY_MARKER_FILE', report_file):
+            dsa._poll_imds()
+        self.assertEqual(m_report_ready.call_count, 1)
+        self.assertTrue(os.path.exists(report_file))
+
+    def test_poll_imds_report_ready_failure_raises_exc_and_doesnt_write_marker(
+            self, m_report_ready, m_request, m_media_switch, m_dhcp, m_net):
+        """poll_imds should write the report_ready marker file if
+        reporting ready succeeds"""
+        report_file = self.tmp_path('report_marker', self.tmp)
+        m_dhcp.return_value = [{
+            'interface': 'eth9', 'fixed-address': '192.168.2.9',
+            'routers': '192.168.2.1', 'subnet-mask': '255.255.255.0',
+            'unknown-245': '624c3620'}]
+        m_media_switch.return_value = None
+        m_report_ready.return_value = False
+        dsa = dsaz.DataSourceAzure({}, distro=None, paths=self.paths)
+        self.assertFalse(os.path.exists(report_file))
+        with mock.patch(MOCKPATH + 'REPORTED_READY_MARKER_FILE', report_file):
+            self.assertRaises(
+                InvalidMetaDataException,
+                dsa._poll_imds)
+        self.assertEqual(m_report_ready.call_count, 1)
+        self.assertFalse(os.path.exists(report_file))
 
 
-@mock.patch(MOCKPATH + 'subp.subp')
-@mock.patch(MOCKPATH + 'util.write_file')
+@mock.patch(MOCKPATH + 'DataSourceAzure._report_ready', mock.MagicMock())
+@mock.patch(MOCKPATH + 'subp.subp', mock.MagicMock())
+@mock.patch(MOCKPATH + 'util.write_file', mock.MagicMock())
 @mock.patch(MOCKPATH + 'util.is_FreeBSD')
 @mock.patch('cloudinit.sources.helpers.netlink.'
             'wait_for_media_disconnect_connect')
@@ -2115,10 +2165,10 @@ class TestAzureDataSourcePreprovisioning(CiTestCase):
         self.paths = helpers.Paths({'cloud_dir': tmp})
         dsaz.BUILTIN_DS_CONFIG['data_dir'] = self.waagent_d
 
-    def test_poll_imds_returns_ovf_env(self, fake_resp,
+    def test_poll_imds_returns_ovf_env(self, m_request,
                                        m_dhcp, m_net,
                                        m_media_switch,
-                                       m_is_bsd, write_f, subp):
+                                       m_is_bsd):
         """The _poll_imds method should return the ovf_env.xml."""
         m_is_bsd.return_value = False
         m_media_switch.return_value = None
@@ -2128,11 +2178,11 @@ class TestAzureDataSourcePreprovisioning(CiTestCase):
         url = 'http://{0}/metadata/reprovisiondata?api-version=2017-04-02'
         host = "169.254.169.254"
         full_url = url.format(host)
-        fake_resp.return_value = mock.MagicMock(status_code=200, text="ovf",
+        m_request.return_value = mock.MagicMock(status_code=200, text="ovf",
                                                 content="ovf")
         dsa = dsaz.DataSourceAzure({}, distro=mock.Mock(), paths=self.paths)
         self.assertTrue(len(dsa._poll_imds()) > 0)
-        self.assertEqual(fake_resp.call_args_list,
+        self.assertEqual(m_request.call_args_list,
                          [mock.call(allow_redirects=True,
                                     headers={'Metadata': 'true',
                                              'User-Agent':
@@ -2147,10 +2197,10 @@ class TestAzureDataSourcePreprovisioning(CiTestCase):
             static_routes=None)
         self.assertEqual(m_net.call_count, 2)
 
-    def test__reprovision_calls__poll_imds(self, fake_resp,
+    def test__reprovision_calls__poll_imds(self, m_request,
                                            m_dhcp, m_net,
                                            m_media_switch,
-                                           m_is_bsd, write_f, subp):
+                                           m_is_bsd):
         """The _reprovision method should call poll IMDS."""
         m_is_bsd.return_value = False
         m_media_switch.return_value = None
@@ -2165,7 +2215,7 @@ class TestAzureDataSourcePreprovisioning(CiTestCase):
         username = "myuser"
         odata = {'HostName': hostname, 'UserName': username}
         content = construct_valid_ovf_env(data=odata)
-        fake_resp.return_value = mock.MagicMock(status_code=200, text=content,
+        m_request.return_value = mock.MagicMock(status_code=200, text=content,
                                                 content=content)
         dsa = dsaz.DataSourceAzure({}, distro=mock.Mock(), paths=self.paths)
         md, _ud, cfg, _d = dsa._reprovision()
@@ -2182,7 +2232,7 @@ class TestAzureDataSourcePreprovisioning(CiTestCase):
                 timeout=dsaz.IMDS_TIMEOUT_IN_SECONDS,
                 url=full_url
             ),
-            fake_resp.call_args_list)
+            m_request.call_args_list)
         self.assertEqual(m_dhcp.call_count, 2)
         m_net.assert_any_call(
             broadcast='192.168.2.255', interface='eth9', ip='192.168.2.9',
diff --git a/tests/unittests/test_datasource/test_azure_helper.py b/tests/unittests/test_datasource/test_azure_helper.py
index 5c31b8be..6e004e34 100644
--- a/tests/unittests/test_datasource/test_azure_helper.py
+++ b/tests/unittests/test_datasource/test_azure_helper.py
@@ -742,7 +742,8 @@ class TestGetMetadataGoalStateXMLAndReportReadyToFabric(CiTestCase):
         self.assertEqual(1, shim.return_value.clean_up.call_count)
 
     @mock.patch.object(azure_helper, 'WALinuxAgentShim')
-    def test_failure_in_registration_calls_clean_up(self, shim):
+    def test_failure_in_registration_propagates_exc_and_calls_clean_up(
+            self, shim):
         shim.return_value.register_with_azure_and_fetch_data.side_effect = (
             SentinelException)
         self.assertRaises(SentinelException,