Repository: Azure/WALinuxAgent
Pull request #635 (2017-03-29 18:46:14)

Pass environment variables to subprocesses
In cases where subprocesses are used to call python processes
it would be helpful to pass on the environment variables used for
the agent itself.  In particular, this is helpful for running
extensions from the agent in system containers.

msftclas: 
@__baude__,
Thanks for your contribution.
To ensure that the project team has proper rights to use your work, please complete the Contribution License Agreement at https://cla.microsoft.com.

It will cover your contributions to all Microsoft-managed open source projects.
_Thanks,
Microsoft Pull Request Bot_
hglkrijger: @baude, would you mind signing the CLA?

Patch:
diff --git a/azurelinuxagent/common/osutil/factory.py b/azurelinuxagent/common/osutil/factory.py
index acd7f6ed..eee9f97a 100644
--- a/azurelinuxagent/common/osutil/factory.py
+++ b/azurelinuxagent/common/osutil/factory.py
@@ -76,6 +76,9 @@ def get_osutil(distro_name=DISTRO_NAME,
         else:
             return RedhatOSUtil()
 
+    elif distro_name == "euleros":
+        return RedhatOSUtil()
+
     elif distro_name == "freebsd":
         return FreeBSDOSUtil()
 
diff --git a/azurelinuxagent/common/protocol/hostplugin.py b/azurelinuxagent/common/protocol/hostplugin.py
index f4ece298..d76326d4 100644
--- a/azurelinuxagent/common/protocol/hostplugin.py
+++ b/azurelinuxagent/common/protocol/hostplugin.py
@@ -38,6 +38,8 @@ MAXIMUM_PAGEBLOB_PAGE_SIZE = 4 * 1024 * 1024  # Max page size: 4MB
 
 
 class HostPluginProtocol(object):
+    _is_default_channel = False
+
     def __init__(self, endpoint, container_id, role_config_name):
         if endpoint is None:
             raise ProtocolError("HostGAPlugin: Endpoint not provided")
@@ -50,6 +52,14 @@ class HostPluginProtocol(object):
         self.role_config_name = role_config_name
         self.manifest_uri = None
 
+    @staticmethod
+    def is_default_channel():
+        return HostPluginProtocol._is_default_channel
+
+    @staticmethod
+    def set_default_channel(is_default):
+        HostPluginProtocol._is_default_channel = is_default
+
     def ensure_initialized(self):
         if not self.is_initialized:
             self.api_versions = self.get_api_versions()
@@ -74,7 +84,7 @@ class HostPluginProtocol(object):
             if response.status != httpclient.OK:
                 logger.error(
                     "HostGAPlugin: Failed Get API versions: {0}".format(
-                        self._read_response_error(response)))
+                        self.read_response_error(response)))
             else:
                 return_val = ustr(remove_bom(response.read()), encoding='utf-8')
 
@@ -165,7 +175,7 @@ class HostPluginProtocol(object):
 
         except Exception as e:
             logger.error("HostGAPlugin: Exception Put VM status: {0}", e)
-    
+
     def _put_block_blob_status(self, sas_url, status_blob):
         url = URI_FORMAT_PUT_VM_STATUS.format(self.endpoint, HOST_PLUGIN_PORT)
 
@@ -178,7 +188,7 @@ class HostPluginProtocol(object):
 
         if response.status != httpclient.OK:
             raise HttpError("HostGAPlugin: Put BlockBlob failed: {0}".format(
-                self._read_response_error(response)))
+                self.read_response_error(response)))
         else:
             logger.verbose("HostGAPlugin: Put BlockBlob status succeeded")
 
@@ -200,7 +210,7 @@ class HostPluginProtocol(object):
         if response.status != httpclient.OK:
             raise HttpError(
                 "HostGAPlugin: Failed PageBlob clean-up: {0}".format(
-                    self._read_response_error(response)))
+                    self.read_response_error(response)))
         else:
             logger.verbose("HostGAPlugin: PageBlob clean-up succeeded")
         
@@ -231,7 +241,7 @@ class HostPluginProtocol(object):
                 raise HttpError(
                     "HostGAPlugin Error: Put PageBlob bytes [{0},{1}]: " \
                     "{2}".format(
-                        start, end, self._read_response_error(response)))
+                        start, end, self.read_response_error(response)))
 
             # Advance to the next page (if any)
             start = end
@@ -266,9 +276,12 @@ class HostPluginProtocol(object):
             return s.decode('utf-8')
         return s
     
-    def _read_response_error(self, response):
+    @staticmethod
+    def read_response_error(response):
+        if response is None:
+            return ''
         body = remove_bom(response.read())
-        if PY_VERSION_MAJOR < 3:
+        if PY_VERSION_MAJOR < 3 and body is not None:
             body = ustr(body, encoding='utf-8')
         return "{0}, {1}, {2}".format(
             response.status,
diff --git a/azurelinuxagent/common/protocol/wire.py b/azurelinuxagent/common/protocol/wire.py
index db343909..051dfa1b 100644
--- a/azurelinuxagent/common/protocol/wire.py
+++ b/azurelinuxagent/common/protocol/wire.py
@@ -370,11 +370,14 @@ class StatusBlob(object):
 
     __storage_version__ = "2014-02-14"
 
+    def prepare(self, blob_type):
+        logger.verbose("Prepare status blob")
+        self.data = self.to_json()
+        self.type = blob_type
+
     def upload(self, url):
         # TODO upload extension only if content has changed
-        logger.verbose("Upload status blob")
         upload_successful = False
-        self.data = self.to_json()
         self.type = self.get_blob_type(url)
         try:
             if self.type == "BlockBlob":
@@ -635,9 +638,14 @@ class WireClient(object):
     def fetch_manifest(self, version_uris):
         logger.verbose("Fetch manifest")
         for version in version_uris:
-            response = self.fetch(version.uri)
+            response = None
+            if not HostPluginProtocol.is_default_channel():
+                response = self.fetch(version.uri)
             if not response:
-                logger.verbose("Manifest could not be downloaded, falling back to host plugin")
+                if HostPluginProtocol.is_default_channel():
+                    logger.verbose("Using host plugin as default channel")
+                else:
+                    logger.verbose("Manifest could not be downloaded, falling back to host plugin")
                 host = self.get_host_plugin()
                 uri, headers = host.get_artifact_request(version.uri)
                 response = self.fetch(uri, headers)
@@ -649,6 +657,9 @@ class WireClient(object):
                 else:
                     host.manifest_uri = version.uri
                     logger.verbose("Manifest downloaded successfully from host plugin")
+                    if not HostPluginProtocol.is_default_channel():
+                        logger.info("Setting host plugin as default channel")
+                        HostPluginProtocol.set_default_channel(True)
             if response:
                 return response
         raise ProtocolError("Failed to fetch manifest from all sources")
@@ -664,12 +675,11 @@ class WireClient(object):
             if resp.status == httpclient.OK:
                 return_value = self.decode_config(resp.read())
             else:
-                logger.warn("Could not fetch {0} [{1}: {2}]",
+                logger.warn("Could not fetch {0} [{1}]",
                             uri,
-                            resp.status,
-                            resp.reason)
+                            HostPluginProtocol.read_response_error(resp))
         except (HttpError, ProtocolError) as e:
-            logger.verbose("Fetch failed from [{0}]", uri)
+            logger.verbose("Fetch failed from [{0}]: {1}", uri, e)
         return return_value
 
     def update_hosting_env(self, goal_state):
@@ -840,10 +850,12 @@ class WireClient(object):
         if ext_conf.status_upload_blob is not None:
             uploaded = False
             try:
-                uploaded = self.status_blob.upload(ext_conf.status_upload_blob)
-                self.report_blob_type(self.status_blob.type,
-                                      ext_conf.status_upload_blob_type)
-            except (HttpError, ProtocolError) as e:
+                self.status_blob.prepare(ext_conf.status_upload_blob_type)
+                if not HostPluginProtocol.is_default_channel():
+                    uploaded = self.status_blob.upload(ext_conf.status_upload_blob)
+                    self.report_blob_type(self.status_blob.type,
+                                          ext_conf.status_upload_blob_type)
+            except (HttpError, ProtocolError):
                 # errors have already been logged
                 pass
             if not uploaded:
@@ -851,6 +863,9 @@ class WireClient(object):
                 host.put_vm_status(self.status_blob,
                                    ext_conf.status_upload_blob,
                                    ext_conf.status_upload_blob_type)
+                if not HostPluginProtocol.is_default_channel():
+                    logger.info("Setting host plugin as default channel")
+                    HostPluginProtocol.set_default_channel(True)
 
     """
     Emit an event to determine if the type in the extension config
diff --git a/azurelinuxagent/common/utils/restutil.py b/azurelinuxagent/common/utils/restutil.py
index 71973705..49d2d68f 100644
--- a/azurelinuxagent/common/utils/restutil.py
+++ b/azurelinuxagent/common/utils/restutil.py
@@ -28,7 +28,7 @@ from azurelinuxagent.common.future import httpclient, urlparse
 REST api util functions
 """
 
-RETRY_WAITING_INTERVAL = 10
+RETRY_WAITING_INTERVAL = 3
 secure_warning = True
 
 
diff --git a/azurelinuxagent/common/version.py b/azurelinuxagent/common/version.py
index 0e3818aa..8ee2c18f 100644
--- a/azurelinuxagent/common/version.py
+++ b/azurelinuxagent/common/version.py
@@ -76,6 +76,9 @@ def get_distro():
         osinfo[2] = "oracle"
         osinfo[3] = "Oracle Linux"
 
+    if os.path.exists("/etc/euleros-release"):
+        osinfo[0] = "euleros"
+
     # The platform.py lib has issue with detecting BIG-IP linux distribution.
     # Merge the following patch provided by F5.
     if os.path.exists("/shared/vadc"):
diff --git a/azurelinuxagent/daemon/resourcedisk/default.py b/azurelinuxagent/daemon/resourcedisk/default.py
index 21de38f0..2b116fb1 100644
--- a/azurelinuxagent/daemon/resourcedisk/default.py
+++ b/azurelinuxagent/daemon/resourcedisk/default.py
@@ -85,6 +85,11 @@ class ResourceDiskHandler(object):
         except ResourceDiskError as e:
             logger.error("Failed to enable swap {0}", e)
 
+    def reread_partition_table(self, device):
+        if shellutil.run("sfdisk -R {0}".format(device), chk_err=False):
+            shellutil.run("blockdev --rereadpt {0}".format(device),
+                          chk_err=False)
+
     def mount_resource_disk(self, mount_point):
         device = self.osutil.device_for_ide_port(1)
         if device is None:
@@ -138,12 +143,13 @@ class ResourceDiskHandler(object):
                 shellutil.run(mkfs_string)
         else:
             logger.info("GPT not detected, determining filesystem")
-            ret = self.change_partition_type(suppress_message=True, option_str="{0} 1".format(device))
+            ret = self.change_partition_type(suppress_message=True, option_str="{0} 1 -n".format(device))
             ptype = ret[1].strip()
             if ptype == "7" and self.fs != "ntfs":
                 logger.info("The partition is formatted with ntfs, updating "
                             "partition type to 83")
                 self.change_partition_type(suppress_message=False, option_str="{0} 1 83".format(device))
+                self.reread_partition_table(device)
                 logger.info("Format partition [{0}]", mkfs_string)
                 shellutil.run(mkfs_string)
             else:
@@ -166,8 +172,8 @@ class ResourceDiskHandler(object):
 
         logger.info("Mount resource disk [{0}]", mount_string)
         ret, output = shellutil.run_get_output(mount_string, chk_err=False)
-        # if the exit code is 32, then the resource disk is already mounted
-        if ret == 32:
+        # if the exit code is 32, then the resource disk can be already mounted
+        if ret == 32 and output.find("is already mounted") != -1:
             logger.warn("Could not mount resource disk: {0}", output)
         elif ret != 0:
             # Some kernels seem to issue an async partition re-read after a
@@ -178,9 +184,7 @@ class ResourceDiskHandler(object):
             logger.warn("Failed to mount resource disk. "
                         "Retry mounting after re-reading partition info.")
 
-            if shellutil.run("sfdisk -R {0}".format(device), chk_err=False):
-                shellutil.run("blockdev --rereadpt {0}".format(device),
-                              chk_err=False)
+            self.reread_partition_table(device)
 
             ret, output = shellutil.run_get_output(mount_string)
             if ret:
diff --git a/azurelinuxagent/ga/exthandlers.py b/azurelinuxagent/ga/exthandlers.py
index 3d3807d3..70bb246e 100644
--- a/azurelinuxagent/ga/exthandlers.py
+++ b/azurelinuxagent/ga/exthandlers.py
@@ -725,7 +725,8 @@ class ExtHandlerInstance(object):
             child = subprocess.Popen(base_dir + "/" + cmd,
                                      shell=True,
                                      cwd=base_dir,
-                                     stdout=devnull)
+                                     stdout=devnull,
+                                     env=os.environ)
         except Exception as e:
             #TODO do not catch all exception
             raise ExtensionError("Failed to launch: {0}, {1}".format(cmd, e))
diff --git a/azurelinuxagent/ga/update.py b/azurelinuxagent/ga/update.py
index 59bc70c5..53f01736 100644
--- a/azurelinuxagent/ga/update.py
+++ b/azurelinuxagent/ga/update.py
@@ -40,6 +40,7 @@ from azurelinuxagent.common.exception import UpdateError, ProtocolError
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.osutil import get_osutil
 from azurelinuxagent.common.protocol import get_protocol_util
+from azurelinuxagent.common.protocol.hostplugin import HostPluginProtocol
 from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
 from azurelinuxagent.common.version import AGENT_NAME, AGENT_VERSION, AGENT_LONG_VERSION, \
                                             AGENT_DIR_GLOB, AGENT_PKG_GLOB, \
@@ -49,7 +50,6 @@ from azurelinuxagent.common.version import AGENT_NAME, AGENT_VERSION, AGENT_LONG
 
 from azurelinuxagent.ga.exthandlers import HandlerManifest
 
-
 AGENT_ERROR_FILE = "error.json" # File name for agent error record
 AGENT_MANIFEST_FILE = "HandlerManifest.json"
 
@@ -140,7 +140,8 @@ class UpdateHandler(object):
                 cmds,
                 cwd=agent_dir,
                 stdout=sys.stdout,
-                stderr=sys.stderr)
+                stderr=sys.stderr,
+                env=os.environ)
 
             logger.info(u"Agent {0} launched with command '{1}'", agent_name, agent_cmd)
 
@@ -698,18 +699,26 @@ class GuestAgent(object):
 
     def _download(self):
         for uri in self.pkg.uris:
-            if self._fetch(uri.uri):
+            if not HostPluginProtocol.is_default_channel() and self._fetch(uri.uri):
                 break
-            else:
-                if self.host is not None and self.host.ensure_initialized():
+            elif self.host is not None and self.host.ensure_initialized():
+                if not HostPluginProtocol.is_default_channel():
                     logger.warn("Download unsuccessful, falling back to host plugin")
-                    uri, headers = self.host.get_artifact_request(uri.uri, self.host.manifest_uri)
-                    if uri is not None \
-                            and headers is not None \
-                            and self._fetch(uri, headers=headers):
-                        break
                 else:
-                    logger.warn("Download unsuccessful, host plugin not available")
+                    logger.verbose("Using host plugin as default channel")
+
+                uri, headers = self.host.get_artifact_request(uri.uri, self.host.manifest_uri)
+                if uri is not None \
+                        and headers is not None \
+                        and self._fetch(uri, headers=headers):
+                    if not HostPluginProtocol.is_default_channel():
+                        logger.verbose("Setting host plugin as default channel")
+                        HostPluginProtocol.set_default_channel(True)
+                    break
+                else:
+                    logger.warn("Host plugin download unsuccessful")
+            else:
+                logger.error("No download channels available")
 
         if not os.path.isfile(self.get_agent_pkg_path()):
             msg = u"Unable to download Agent {0} from any URI".format(self.name)
@@ -732,6 +741,9 @@ class GuestAgent(object):
                                     bytearray(package),
                                     asbin=True)
                 logger.info(u"Agent {0} downloaded from {1}", self.name, uri)
+            else:
+                logger.verbose("Fetch was unsuccessful [{0}]",
+                               HostPluginProtocol.read_response_error(resp))
         except restutil.HttpError as http_error:
             logger.verbose(u"Agent {0} download from {1} failed [{2}]",
                            self.name,
diff --git a/azurelinuxagent/pa/rdma/centos.py b/azurelinuxagent/pa/rdma/centos.py
index 214f9ead..942d225f 100644
--- a/azurelinuxagent/pa/rdma/centos.py
+++ b/azurelinuxagent/pa/rdma/centos.py
@@ -82,7 +82,7 @@ class CentOSRDMAHandler(RDMAHandler):
         # Example match (pkg name, -, followed by 3 segments, fw_version and -):
         # - pkg=microsoft-hyper-v-rdma-4.1.0.142-20160323.x86_64
         # - fw_version=142
-        pattern = '{0}-\d\.\d\.\d\.({1})-'.format(
+        pattern = '{0}-(\d+\.){3,}({1})-'.format(
             self.rdma_user_mode_package_name, fw_version)
         return re.match(pattern, pkg)
 
@@ -156,7 +156,7 @@ class CentOSRDMAHandler(RDMAHandler):
 
         # Install kernel mode driver (kmod-microsoft-hyper-v-rdma-*)
         kmod_pkg = self.get_file_by_pattern(
-            pkgs, "%s-\d\.\d\.\d\.+(%s)-\d{8}\.x86_64.rpm" % (self.rdma_kernel_mode_package_name, fw_version))
+            pkgs, "%s-(\d+\.){3,}(%s)-\d{8}\.x86_64.rpm" % (self.rdma_kernel_mode_package_name, fw_version))
         if not kmod_pkg:
             raise Exception("RDMA kernel mode package not found")
         kmod_pkg_path = os.path.join(pkg_dir, kmod_pkg)
@@ -165,7 +165,7 @@ class CentOSRDMAHandler(RDMAHandler):
 
         # Install user mode driver (microsoft-hyper-v-rdma-*)
         umod_pkg = self.get_file_by_pattern(
-            pkgs, "%s-\d\.\d\.\d\.+(%s)-\d{8}\.x86_64.rpm" % (self.rdma_user_mode_package_name, fw_version))
+            pkgs, "%s-(\d+\.){3,}(%s)-\d{8}\.x86_64.rpm" % (self.rdma_user_mode_package_name, fw_version))
         if not umod_pkg:
             raise Exception("RDMA user mode package not found")
         umod_pkg_path = os.path.join(pkg_dir, umod_pkg)
diff --git a/bin/waagent2.0 b/bin/waagent2.0
index 80af1c73..1a72ba73 100644
--- a/bin/waagent2.0
+++ b/bin/waagent2.0
@@ -1093,6 +1093,18 @@ class centosDistro(redhatDistro):
     def __init__(self):
         super(centosDistro,self).__init__()
 
+############################################################    
+#   eulerosDistro
+############################################################    
+
+class eulerosDistro(redhatDistro):
+    """
+    EulerOS Distro concrete class
+    Put EulerOS specific behavior here...
+    """
+    def __init__(self):
+        super(eulerosDistro,self).__init__()
+
 ############################################################    
 #	oracleDistro
 ############################################################    
@@ -6016,6 +6028,8 @@ def DistInfo(fullname=0):
     if 'linux_distribution' in dir(platform):
         distinfo = list(platform.linux_distribution(full_distribution_name=fullname))
         distinfo[0] = distinfo[0].strip() # remove trailing whitespace in distro name
+        if os.path.exists("/etc/euleros-release"):
+            distinfo[0] = "euleros"
         return distinfo
     else:
         return platform.dist()
diff --git a/config/66-azure-storage.rules b/config/66-azure-storage.rules
index ab306288..5b2b799f 100644
--- a/config/66-azure-storage.rules
+++ b/config/66-azure-storage.rules
@@ -1,12 +1,22 @@
-ACTION!="add|change", GOTO="azure_end"
-SUBSYSTEM!="block", GOTO="azure_end"
-ATTRS{ID_VENDOR}!="Msft", GOTO="azure_end"
-ATTRS{ID_MODEL}!="Virtual_Disk", GOTO="azure_end"
+ACTION=="add|change", SUBSYSTEM=="block", ENV{ID_VENDOR}=="Msft", ENV{ID_MODEL}=="Virtual_Disk", GOTO="azure_disk"
+GOTO="azure_end"
 
+LABEL="azure_disk"
 # Root has a GUID of 0000 as the second value
 # The resource/resource has GUID of 0001 as the second value
 ATTRS{device_id}=="?00000000-0000-*", ENV{fabric_name}="root", GOTO="azure_names"
 ATTRS{device_id}=="?00000000-0001-*", ENV{fabric_name}="resource", GOTO="azure_names"
+# Wellknown SCSI controllers
+ATTRS{device_id}=="{f8b3781a-1e82-4818-a1c3-63d806ec15bb}", ENV{fabric_scsi_controller}="scsi0", GOTO="azure_datadisk"
+ATTRS{device_id}=="{f8b3781b-1e82-4818-a1c3-63d806ec15bb}", ENV{fabric_scsi_controller}="scsi1", GOTO="azure_datadisk"
+ATTRS{device_id}=="{f8b3781c-1e82-4818-a1c3-63d806ec15bb}", ENV{fabric_scsi_controller}="scsi2", GOTO="azure_datadisk"
+ATTRS{device_id}=="{f8b3781d-1e82-4818-a1c3-63d806ec15bb}", ENV{fabric_scsi_controller}="scsi3", GOTO="azure_datadisk"
+GOTO="azure_end"
+
+# Retrieve LUN number for datadisks
+LABEL="azure_datadisk"
+ENV{DEVTYPE}=="partition", PROGRAM="/bin/sh -c 'readlink /sys/class/block/%k/../device|cut -d: -f4'", ENV{fabric_name}="$env{fabric_scsi_controller}/lun$result", GOTO="azure_names"
+PROGRAM="/bin/sh -c 'readlink /sys/class/block/%k/device|cut -d: -f4'", ENV{fabric_name}="$env{fabric_scsi_controller}/lun$result", GOTO="azure_names"
 GOTO="azure_end"
 
 # Create the symlinks
@@ -15,4 +25,3 @@ ENV{DEVTYPE}=="disk", SYMLINK+="disk/azure/$env{fabric_name}"
 ENV{DEVTYPE}=="partition", SYMLINK+="disk/azure/$env{fabric_name}-part%n"
 
 LABEL="azure_end"
-


Test patch:
diff --git a/tests/protocol/test_hostplugin.py b/tests/protocol/test_hostplugin.py
index 078af706..8ebb468c 100644
--- a/tests/protocol/test_hostplugin.py
+++ b/tests/protocol/test_hostplugin.py
@@ -113,21 +113,24 @@ class TestHostPlugin(AgentTestCase):
 
     def test_fallback(self):
         """
-        Validate fallback to upload status using HostGAPlugin is happening when status reporting via
-        default method is unsuccessful
+        Validate fallback to upload status using HostGAPlugin is happening when
+        status reporting via default method is unsuccessful
         """
         test_goal_state = wire.GoalState(WireProtocolData(DATA_FILE).goal_state)
-
+        status = restapi.VMStatus(status="Ready", message="Guest Agent is running")
         with patch.object(wire.HostPluginProtocol, "put_vm_status") as patch_put:
             with patch.object(wire.StatusBlob, "upload", return_value=False) as patch_upload:
                 wire_protocol_client = wire.WireProtocol(wireserver_url).client
                 wire_protocol_client.get_goal_state = Mock(return_value=test_goal_state)
                 wire_protocol_client.ext_conf = wire.ExtensionsConfig(None)
                 wire_protocol_client.ext_conf.status_upload_blob = sas_url
+                wire_protocol_client.status_blob.set_vm_status(status)
                 wire_protocol_client.upload_status_blob()
                 self.assertTrue(patch_put.call_count == 1,
                                 "Fallback was not engaged")
                 self.assertTrue(patch_put.call_args[0][1] == sas_url)
+                self.assertTrue(wire.HostPluginProtocol.is_default_channel())
+                wire.HostPluginProtocol.set_default_channel(False)
 
     def test_validate_http_request(self):
         """Validate correct set of data is sent to HostGAPlugin when reporting VM status"""
@@ -161,19 +164,19 @@ class TestHostPlugin(AgentTestCase):
 
     def test_no_fallback(self):
         """
-        Validate fallback to upload status using HostGAPlugin is not happening when status reporting via
-        default method is successful
+        Validate fallback to upload status using HostGAPlugin is not happening
+        when status reporting via default method is successful
         """
-        with patch.object(wire.HostPluginProtocol,
-                          "put_vm_status") as patch_put:
+        vmstatus = restapi.VMStatus(message="Ready", status="Ready")
+        with patch.object(wire.HostPluginProtocol, "put_vm_status") as patch_put:
             with patch.object(wire.StatusBlob, "upload") as patch_upload:
                 patch_upload.return_value = True
                 wire_protocol_client = wire.WireProtocol(wireserver_url).client
                 wire_protocol_client.ext_conf = wire.ExtensionsConfig(None)
                 wire_protocol_client.ext_conf.status_upload_blob = sas_url
+                wire_protocol_client.status_blob.vm_status = vmstatus
                 wire_protocol_client.upload_status_blob()
-                self.assertTrue(patch_put.call_count == 0,
-                                "Fallback was engaged")
+                self.assertTrue(patch_put.call_count == 0, "Fallback was engaged")
 
     def test_validate_block_blob(self):
         """Validate correct set of data is sent to HostGAPlugin when reporting VM status"""
diff --git a/tests/protocol/test_wire.py b/tests/protocol/test_wire.py
index e083678b..e36fc652 100644
--- a/tests/protocol/test_wire.py
+++ b/tests/protocol/test_wire.py
@@ -148,13 +148,15 @@ class TestWireProtocolGetters(AgentTestCase):
                                      host_uri)
 
     def test_upload_status_blob_default(self, *args):
+        vmstatus = VMStatus(message="Ready", status="Ready")
         wire_protocol_client = WireProtocol(wireserver_url).client
         wire_protocol_client.ext_conf = ExtensionsConfig(None)
         wire_protocol_client.ext_conf.status_upload_blob = testurl
+        wire_protocol_client.status_blob.vm_status = vmstatus
 
         with patch.object(WireClient, "get_goal_state") as patch_get_goal_state:
             with patch.object(HostPluginProtocol, "put_vm_status") as patch_host_ga_plugin_upload:
-                with patch.object(StatusBlob, "upload", return_value = True) as patch_default_upload:
+                with patch.object(StatusBlob, "upload", return_value=True) as patch_default_upload:
                     wire_protocol_client.upload_status_blob()
 
                     patch_default_upload.assert_called_once_with(testurl)
@@ -162,20 +164,24 @@ class TestWireProtocolGetters(AgentTestCase):
                     patch_host_ga_plugin_upload.assert_not_called()
 
     def test_upload_status_blob_host_ga_plugin(self, *args):
+        vmstatus = VMStatus(message="Ready", status="Ready")
         wire_protocol_client = WireProtocol(wireserver_url).client
         wire_protocol_client.ext_conf = ExtensionsConfig(None)
         wire_protocol_client.ext_conf.status_upload_blob = testurl
         wire_protocol_client.ext_conf.status_upload_blob_type = testtype
+        wire_protocol_client.status_blob.vm_status = vmstatus
         goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
 
         with patch.object(HostPluginProtocol, "put_vm_status") as patch_host_ga_plugin_upload:
             with patch.object(StatusBlob, "upload", return_value=False) as patch_default_upload:
-                wire_protocol_client.get_goal_state = Mock(return_value = goal_state)
+                wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
                 wire_protocol_client.upload_status_blob()
 
                 patch_default_upload.assert_called_once_with(testurl)
                 wire_protocol_client.get_goal_state.assert_called_once()
                 patch_host_ga_plugin_upload.assert_called_once_with(wire_protocol_client.status_blob, testurl, testtype)
+                self.assertTrue(HostPluginProtocol.is_default_channel())
+                HostPluginProtocol.set_default_channel(False)
 
     def test_get_in_vm_artifacts_profile_blob_not_available(self, *args):
         wire_protocol_client = WireProtocol(wireserver_url).client