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

1 files changed, 46 insertions, 84 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