sources/azure: refactor _should_reprovision[_after_nic_attach]() logic (#1206)

Consolidate _should_reprovision_after_nic_attach() with
_should_reprovision() into the following:

_write_reprovision_marker() to write provisioning marker for
reboot-during-provisioning case.

PPSType enum and _determine_pps_type() for determining which to
provisioning mode, if any, we're running under.

PPSType.UNKNOWN is when the reprovisioning marker is found and we
do not have the context to know what the original mode was. In this
scenario, we must resort to polling for reprovision data.

Tests:

Introduce a simple data source fixture to for fine-grain
control of mocking with pytest without unittest.

Migrate relevant _should_reprovision() tests into a combination of
TestDeterminePPSTypeScenarios cases.

Signed-off-by: Chris Patterson cpatterson@microsoft.com

2 files changed, 135 insertions, 167 deletions

diff --git a/cloudinit/sources/DataSourceAzure.py b/cloudinit/sources/DataSourceAzure.py
index 4e88c22b..3cd74a63 100755
--- a/cloudinit/sources/DataSourceAzure.py
+++ b/cloudinit/sources/DataSourceAzure.py
@@ -55,7 +55,6 @@ DEFAULT_FS = "ext4"
 # DMI chassis-asset-tag is set static for all azure instances
 AZURE_CHASSIS_ASSET_TAG = "7783-7084-3265-9085-8269-3286-77"
 REPROVISION_MARKER_FILE = "/var/lib/cloud/data/poll_imds"
-REPROVISION_NIC_ATTACH_MARKER_FILE = "/var/lib/cloud/data/wait_for_nic_attach"
 REPROVISION_NIC_DETACHED_MARKER_FILE = "/var/lib/cloud/data/nic_detached"
 REPORTED_READY_MARKER_FILE = "/var/lib/cloud/data/reported_ready"
 AGENT_SEED_DIR = "/var/lib/waagent"
@@ -80,6 +79,13 @@ class MetadataType(Enum):
     REPROVISION_DATA = "{}/reprovisiondata".format(IMDS_URL)
 
 
+class PPSType(Enum):
+    NONE = "None"
+    RUNNING = "Running"
+    SAVABLE = "Savable"
+    UNKNOWN = "Unknown"
+
+
 PLATFORM_ENTROPY_SOURCE = "/sys/firmware/acpi/tables/OEM0"
 
 # List of static scripts and network config artifacts created by
@@ -337,31 +343,19 @@ class DataSourceAzure(sources.DataSource):
         # it determines the value of ret. More specifically, the first one in
         # the candidate list determines the path to take in order to get the
         # metadata we need.
-        reprovision = False
         ovf_is_accessible = False
-        reprovision_after_nic_attach = False
         metadata_source = None
         md = {}
         userdata_raw = ""
         cfg = {}
         files = {}
         if os.path.isfile(REPROVISION_MARKER_FILE):
-            reprovision = True
             metadata_source = "IMDS"
             report_diagnostic_event(
                 "Reprovision marker file already present "
                 "before crawling Azure metadata: %s" % REPROVISION_MARKER_FILE,
                 logger_func=LOG.debug,
             )
-        elif os.path.isfile(REPROVISION_NIC_ATTACH_MARKER_FILE):
-            reprovision_after_nic_attach = True
-            metadata_source = "NIC_ATTACH_MARKER_PRESENT"
-            report_diagnostic_event(
-                "Reprovision nic attach marker file "
-                "already present before crawling Azure "
-                "metadata: %s" % REPROVISION_NIC_ATTACH_MARKER_FILE,
-                logger_func=LOG.debug,
-            )
         else:
             for src in list_possible_azure_ds(self.seed_dir, ddir):
                 try:
@@ -418,21 +412,18 @@ class DataSourceAzure(sources.DataSource):
             report_diagnostic_event(msg)
             raise sources.InvalidMetaDataException(msg)
 
-        perform_reprovision = reprovision or self._should_reprovision(
-            cfg, imds_md
-        )
-        perform_reprovision_after_nic_attach = (
-            reprovision_after_nic_attach
-            or self._should_reprovision_after_nic_attach(cfg, imds_md)
-        )
-
-        if perform_reprovision or perform_reprovision_after_nic_attach:
+        pps_type = self._determine_pps_type(cfg, imds_md)
+        if pps_type != PPSType.NONE:
             if util.is_FreeBSD():
                 msg = "Free BSD is not supported for PPS VMs"
                 report_diagnostic_event(msg, logger_func=LOG.error)
                 raise sources.InvalidMetaDataException(msg)
-            if perform_reprovision_after_nic_attach:
+
+            self._write_reprovision_marker()
+
+            if pps_type == PPSType.SAVABLE:
                 self._wait_for_all_nics_ready()
+
             md, userdata_raw, cfg, files = self._reprovision()
             # fetch metadata again as it has changed after reprovisioning
             imds_md = self.get_imds_data_with_api_fallback(
@@ -515,7 +506,7 @@ class DataSourceAzure(sources.DataSource):
             crawled_data["metadata"]["random_seed"] = seed
         crawled_data["metadata"]["instance-id"] = self._iid()
 
-        if perform_reprovision or perform_reprovision_after_nic_attach:
+        if pps_type != PPSType.NONE:
             LOG.info("Reporting ready to Azure after getting ReprovisionData")
             use_cached_ephemeral = (
                 self.distro.networking.is_up(self.fallback_interface)
@@ -1398,67 +1389,39 @@ class DataSourceAzure(sources.DataSource):
             )
             return None
 
-    def _should_reprovision_after_nic_attach(
-        self, cfg: dict, imds_md=None
-    ) -> bool:
-        """Whether or not we should wait for nic attach and then poll
-        IMDS for reprovisioning data. Also sets a marker file to poll IMDS.
-
-        The marker file is used for the following scenario: the VM boots into
-        wait for nic attach, which we expect to be proceeding infinitely until
-        the nic is attached. If for whatever reason the platform moves us to a
-        new host (for instance a hardware issue), we need to keep waiting.
-        However, since the VM reports ready to the Fabric, we will not attach
-        the ISO, thus cloud-init needs to have a way of knowing that it should
-        jump back into the waiting mode in order to retrieve the ovf_env.
-
-        @param cfg: OVF cfg.
-        @param imds_md: Metadata obtained from IMDS
-        @return: Whether to reprovision after waiting for nics to be attached.
-        """
-        path = REPROVISION_NIC_ATTACH_MARKER_FILE
-        if (
-            cfg.get("PreprovisionedVMType", None) == "Savable"
-            or self._ppstype_from_imds(imds_md) == "Savable"
-            or os.path.isfile(path)
+    def _determine_pps_type(self, ovf_cfg: dict, imds_md: dict) -> PPSType:
+        """Determine PPS type using OVF, IMDS data, and reprovision marker."""
+        if os.path.isfile(REPROVISION_MARKER_FILE):
+            pps_type = PPSType.UNKNOWN
+        elif (
+            ovf_cfg.get("PreprovisionedVMType", None) == PPSType.SAVABLE.value
+            or self._ppstype_from_imds(imds_md) == PPSType.SAVABLE.value
         ):
-            if not os.path.isfile(path):
-                LOG.info(
-                    "Creating a marker file to wait for nic attach: %s", path
-                )
-                util.write_file(
-                    path,
-                    "{pid}: {time}\n".format(pid=os.getpid(), time=time()),
-                )
-            return True
-        return False
-
-    def _should_reprovision(self, cfg: dict, imds_md=None):
-        """Whether or not we should poll IMDS for reprovisioning data.
-        Also sets a marker file to poll IMDS.
-
-        The marker file is used for the following scenario: the VM boots into
-        this polling loop, which we expect to be proceeding infinitely until
-        the VM is picked. If for whatever reason the platform moves us to a
-        new host (for instance a hardware issue), we need to keep polling.
-        However, since the VM reports ready to the Fabric, we will not attach
-        the ISO, thus cloud-init needs to have a way of knowing that it should
-        jump back into the polling loop in order to retrieve the ovf_env."""
-        path = REPROVISION_MARKER_FILE
-        if (
-            cfg.get("PreprovisionedVm") is True
-            or cfg.get("PreprovisionedVMType", None) == "Running"
-            or self._ppstype_from_imds(imds_md) == "Running"
-            or os.path.isfile(path)
+            pps_type = PPSType.SAVABLE
+        elif (
+            ovf_cfg.get("PreprovisionedVm") is True
+            or ovf_cfg.get("PreprovisionedVMType", None)
+            == PPSType.RUNNING.value
+            or self._ppstype_from_imds(imds_md) == PPSType.RUNNING.value
         ):
-            if not os.path.isfile(path):
-                LOG.info("Creating a marker file to poll imds: %s", path)
-                util.write_file(
-                    path,
-                    "{pid}: {time}\n".format(pid=os.getpid(), time=time()),
-                )
-            return True
-        return False
+            pps_type = PPSType.RUNNING
+        else:
+            pps_type = PPSType.NONE
+
+        report_diagnostic_event(
+            "PPS type: %s" % pps_type.value, logger_func=LOG.info
+        )
+        return pps_type
+
+    def _write_reprovision_marker(self):
+        """Write reprovision marker file in case system is rebooted."""
+        LOG.info(
+            "Creating a marker file to poll imds: %s", REPROVISION_MARKER_FILE
+        )
+        util.write_file(
+            REPROVISION_MARKER_FILE,
+            "{pid}: {time}\n".format(pid=os.getpid(), time=time()),
+        )
 
     def _reprovision(self):
         """Initiate the reprovisioning workflow."""
@@ -1507,7 +1470,6 @@ class DataSourceAzure(sources.DataSource):
 
         util.del_file(REPORTED_READY_MARKER_FILE)
         util.del_file(REPROVISION_MARKER_FILE)
-        util.del_file(REPROVISION_NIC_ATTACH_MARKER_FILE)
         util.del_file(REPROVISION_NIC_DETACHED_MARKER_FILE)
         return fabric_data
 
diff --git a/tests/unittests/sources/test_azure.py b/tests/unittests/sources/test_azure.py
index 66169a7a..5581492e 100644
--- a/tests/unittests/sources/test_azure.py
+++ b/tests/unittests/sources/test_azure.py
@@ -8,6 +8,7 @@ import stat
 import xml.etree.ElementTree as ET
 
 import httpretty
+import pytest
 import requests
 import yaml
 
@@ -37,6 +38,13 @@ from tests.unittests.helpers import (
 )
 
 
+@pytest.fixture
+def azure_ds(request, paths):
+    """Provide DataSourceAzure instance with mocks for minimal test case."""
+    with mock.patch(MOCKPATH + "_is_platform_viable", return_value=True):
+        yield dsaz.DataSourceAzure(sys_cfg={}, distro=mock.Mock(), paths=paths)
+
+
 def construct_valid_ovf_env(
     data=None, pubkeys=None, userdata=None, platform_settings=None
 ):
@@ -1243,40 +1251,6 @@ scbus-1 on xpt0 bus 0
 
     @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,
-    ):
-        """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"
     )
     @mock.patch(
@@ -2761,61 +2735,93 @@ class TestPreprovisioningReadAzureOvfFlag(CiTestCase):
         self.assertEqual("Savable", cfg["PreprovisionedVMType"])
 
 
-@mock.patch("os.path.isfile")
-class TestPreprovisioningShouldReprovision(CiTestCase):
+@pytest.mark.parametrize(
+    "ovf_cfg,imds_md,pps_type",
+    [
+        (
+            {"PreprovisionedVm": False, "PreprovisionedVMType": None},
+            {},
+            dsaz.PPSType.NONE,
+        ),
+        (
+            {"PreprovisionedVm": True, "PreprovisionedVMType": "Running"},
+            {},
+            dsaz.PPSType.RUNNING,
+        ),
+        (
+            {"PreprovisionedVm": True, "PreprovisionedVMType": "Savable"},
+            {},
+            dsaz.PPSType.SAVABLE,
+        ),
+        (
+            {"PreprovisionedVm": True},
+            {},
+            dsaz.PPSType.RUNNING,
+        ),
+        (
+            {},
+            {"extended": {"compute": {"ppsType": "None"}}},
+            dsaz.PPSType.NONE,
+        ),
+        (
+            {},
+            {"extended": {"compute": {"ppsType": "Running"}}},
+            dsaz.PPSType.RUNNING,
+        ),
+        (
+            {},
+            {"extended": {"compute": {"ppsType": "Savable"}}},
+            dsaz.PPSType.SAVABLE,
+        ),
+        (
+            {"PreprovisionedVm": False, "PreprovisionedVMType": None},
+            {"extended": {"compute": {"ppsType": "None"}}},
+            dsaz.PPSType.NONE,
+        ),
+        (
+            {"PreprovisionedVm": True, "PreprovisionedVMType": "Running"},
+            {"extended": {"compute": {"ppsType": "Running"}}},
+            dsaz.PPSType.RUNNING,
+        ),
+        (
+            {"PreprovisionedVm": True, "PreprovisionedVMType": "Savable"},
+            {"extended": {"compute": {"ppsType": "Savable"}}},
+            dsaz.PPSType.SAVABLE,
+        ),
+        (
+            {"PreprovisionedVm": True},
+            {"extended": {"compute": {"ppsType": "Running"}}},
+            dsaz.PPSType.RUNNING,
+        ),
+    ],
+)
+class TestDeterminePPSTypeScenarios:
+    @mock.patch("os.path.isfile", return_value=False)
+    def test_determine_pps_without_reprovision_marker(
+        self, is_file, azure_ds, ovf_cfg, imds_md, pps_type
+    ):
+        assert azure_ds._determine_pps_type(ovf_cfg, imds_md) == pps_type
+
+    @mock.patch("os.path.isfile", return_value=True)
+    def test_determine_pps_with_reprovision_marker(
+        self, is_file, azure_ds, ovf_cfg, imds_md, pps_type
+    ):
+        assert (
+            azure_ds._determine_pps_type(ovf_cfg, imds_md)
+            == dsaz.PPSType.UNKNOWN
+        )
+        assert is_file.mock_calls == [mock.call(dsaz.REPROVISION_MARKER_FILE)]
+
+
+@mock.patch("os.path.isfile", return_value=False)
+class TestReprovision(CiTestCase):
     def setUp(self):
-        super(TestPreprovisioningShouldReprovision, self).setUp()
+        super(TestReprovision, self).setUp()
         tmp = self.tmp_dir()
         self.waagent_d = self.tmp_path("/var/lib/waagent", tmp)
         self.paths = helpers.Paths({"cloud_dir": tmp})
         dsaz.BUILTIN_DS_CONFIG["data_dir"] = self.waagent_d
 
-    @mock.patch(MOCKPATH + "util.write_file")
-    def test__should_reprovision_with_true_cfg(self, isfile, write_f):
-        """The _should_reprovision method should return true with config
-        flag present."""
-        isfile.return_value = False
-        dsa = dsaz.DataSourceAzure({}, distro=mock.Mock(), paths=self.paths)
-        self.assertTrue(
-            dsa._should_reprovision({"PreprovisionedVm": True}, None)
-        )
-
-    def test__should_reprovision_with_file_existing(self, isfile):
-        """The _should_reprovision method should return True if the sentinal
-        exists."""
-        isfile.return_value = True
-        dsa = dsaz.DataSourceAzure({}, distro=mock.Mock(), paths=self.paths)
-        self.assertTrue(
-            dsa._should_reprovision({"preprovisionedvm": False}, None)
-        )
-
-    def test__should_reprovision_returns_false(self, isfile):
-        """The _should_reprovision method should return False
-        if config and sentinal are not present."""
-        isfile.return_value = False
-        dsa = dsaz.DataSourceAzure({}, distro=mock.Mock(), paths=self.paths)
-        self.assertFalse(dsa._should_reprovision({}))
-
-    @mock.patch(MOCKPATH + "util.write_file", autospec=True)
-    def test__should_reprovision_uses_imds_md(self, write_file, isfile):
-        """The _should_reprovision method should be able to
-        retrieve the preprovisioning VM type from imds metadata"""
-        isfile.return_value = False
-        dsa = dsaz.DataSourceAzure({}, distro=mock.Mock(), paths=self.paths)
-        self.assertTrue(
-            dsa._should_reprovision(
-                {},
-                {"extended": {"compute": {"ppsType": "Running"}}},
-            )
-        )
-        self.assertFalse(dsa._should_reprovision({}, {}))
-        self.assertFalse(
-            dsa._should_reprovision(
-                {},
-                {"extended": {"compute": {"hasCustomData": False}}},
-            )
-        )
-
     @mock.patch(MOCKPATH + "DataSourceAzure._poll_imds")
     def test_reprovision_calls__poll_imds(self, _poll_imds, isfile):
         """_reprovision will poll IMDS."""