Repository: Azure/WALinuxAgent
Pull request #453 (2016-09-29 01:01:19)

PR for #452, Add Overprovisioning support


Patch:
diff --git a/azurelinuxagent/common/protocol/hostplugin.py b/azurelinuxagent/common/protocol/hostplugin.py
index 66d58898..33ed2ec3 100644
--- a/azurelinuxagent/common/protocol/hostplugin.py
+++ b/azurelinuxagent/common/protocol/hostplugin.py
@@ -22,19 +22,25 @@ from azurelinuxagent.common.utils import textutil
 
 HOST_PLUGIN_PORT = 32526
 URI_FORMAT_GET_API_VERSIONS = "http://{0}:{1}/versions"
+URI_FORMAT_GET_EXTENSION_ARTIFACT = "http://{0}:{1}/extensionArtifact"
 URI_FORMAT_PUT_VM_STATUS = "http://{0}:{1}/status"
 URI_FORMAT_PUT_LOG = "http://{0}:{1}/vmAgentLog"
 API_VERSION = "2015-09-01"
+HEADER_CONTAINER_ID = "x-ms-containerid"
+HEADER_VERSION = "x-ms-version"
+HEADER_HOST_CONFIG_NAME = "x-ms-host-config-name"
 
 
 class HostPluginProtocol(object):
-    def __init__(self, endpoint):
+    def __init__(self, endpoint, container_id, role_config_name):
         if endpoint is None:
             raise ProtocolError("Host plugin endpoint not provided")
         self.is_initialized = False
         self.is_available = False
         self.api_versions = None
         self.endpoint = endpoint
+        self.container_id = container_id
+        self.role_config_name = role_config_name
 
     def ensure_initialized(self):
         if not self.is_initialized:
@@ -48,7 +54,8 @@ class HostPluginProtocol(object):
                                                  HOST_PLUGIN_PORT)
         logger.info("getting API versions at [{0}]".format(url))
         try:
-            response = restutil.http_get(url)
+            headers = { HEADER_CONTAINER_ID: self.container_id }
+            response = restutil.http_get(url, headers)
             if response.status != httpclient.OK:
                 logger.error(
                     "get API versions returned status code [{0}]".format(
@@ -59,6 +66,24 @@ class HostPluginProtocol(object):
             logger.error("get API versions failed with [{0}]".format(e))
             return []
 
+    def get_extension_artifact_url_and_headers(self, artifact_url, artifact_manifest_url = None):
+        if not self.ensure_initialized():
+            logger.error("host plugin channel is not available")
+            return
+        if artifact_url is None or artifact_url.isspace():
+            logger.error("no extension artifact url was provided")
+            return
+
+        url = URI_FORMAT_GET_EXTENSION_ARTIFACT.format(self.endpoint, HOST_PLUGIN_PORT)
+        headers = {HEADER_VERSION: API_VERSION,
+                   HEADER_CONTAINER_ID: self.container_id,
+                   HEADER_HOST_CONFIG_NAME: self.role_config_name,
+                   "x-ms-artifact-location": artifact_url}
+        if artifact_manifest_url is not None:
+            headers["x-ms-artifact-manifest-location"] = artifact_manifest_url
+
+        return url, headers
+
     def put_vm_status(self, status_blob, sas_url):
         """
         Try to upload the VM status via the host plugin /status channel
@@ -73,8 +98,11 @@ class HostPluginProtocol(object):
             return
         url = URI_FORMAT_PUT_VM_STATUS.format(self.endpoint, HOST_PLUGIN_PORT)
         logger.verbose("Posting VM status to host channel [{0}]".format(url))
-        status = textutil.b64encode(status_blob.vm_status.vmAgent.status)
-        headers = {"x-ms-version": API_VERSION}
+        status = textutil.b64encode(status_blob.data)
+        headers = {HEADER_VERSION: API_VERSION,
+                   "Content-type": "application/json",
+                   HEADER_CONTAINER_ID: self.container_id,
+                   HEADER_HOST_CONFIG_NAME: self.role_config_name}
         blob_headers = [{'headerName': 'x-ms-version',
                          'headerValue': status_blob.__storage_version__},
                         {'headerName': 'x-ms-blob-type',
@@ -89,7 +117,7 @@ class HostPluginProtocol(object):
         except HttpError as e:
             logger.error("put VM status failed with [{0}]".format(e))
 
-    def put_vm_log(self, content, container_id, deployment_id):
+    def put_vm_log(self, content):
         """
         Try to upload the given content to the host plugin
         :param deployment_id: the deployment id, which is obtained from the
@@ -102,18 +130,18 @@ class HostPluginProtocol(object):
         if not self.ensure_initialized():
             logger.error("host plugin channel is not available")
             return
-        if content is None or container_id is None or deployment_id is None:
+        if content is None or self.goal_state.container_id is None or self.goal_state.deployment_id is None:
             logger.error(
                 "invalid arguments passed: "
                 "[{0}], [{1}], [{2}]".format(
                     content,
-                    container_id,
-                    deployment_id))
+                    self.goal_state.container_id,
+                    self.goal_state.deployment_id))
             return
         url = URI_FORMAT_PUT_LOG.format(self.endpoint, HOST_PLUGIN_PORT)
 
-        headers = {"x-ms-vmagentlog-deploymentid": deployment_id,
-                   "x-ms-vmagentlog-containerid": container_id}
+        headers = {"x-ms-vmagentlog-deploymentid": self.goal_state.deployment_id,
+                   "x-ms-vmagentlog-containerid": self.goal_state.container_id}
         logger.info("put VM log at [{0}]".format(url))
         try:
             response = restutil.http_put(url, content, headers)
diff --git a/azurelinuxagent/common/protocol/restapi.py b/azurelinuxagent/common/protocol/restapi.py
index aa7872bb..d9ade02f 100644
--- a/azurelinuxagent/common/protocol/restapi.py
+++ b/azurelinuxagent/common/protocol/restapi.py
@@ -282,6 +282,9 @@ class Protocol(DataContract):
     def get_ext_handler_pkgs(self, extension):
         raise NotImplementedError()
 
+    def get_in_vm_artifacts_profile(self):
+        raise NotImplementedError()
+
     def download_ext_handler_pkg(self, uri):
         try:
             resp = restutil.http_get(uri, chk_proxy=True)
diff --git a/azurelinuxagent/common/protocol/wire.py b/azurelinuxagent/common/protocol/wire.py
index a958bc8d..a9c2e02f 100644
--- a/azurelinuxagent/common/protocol/wire.py
+++ b/azurelinuxagent/common/protocol/wire.py
@@ -25,7 +25,7 @@ import azurelinuxagent.common.conf as conf
 from azurelinuxagent.common.exception import ProtocolNotFoundError
 from azurelinuxagent.common.future import httpclient, bytebuffer
 from azurelinuxagent.common.utils.textutil import parse_doc, findall, find, \
-    findtext, getattrib, gettext, remove_bom, get_bytes_from_pem
+    findtext, getattrib, gettext, remove_bom, get_bytes_from_pem, parse_json
 import azurelinuxagent.common.utils.fileutil as fileutil
 from azurelinuxagent.common.utils.cryptutil import CryptUtil
 from azurelinuxagent.common.protocol.restapi import *
@@ -139,6 +139,10 @@ class WireProtocol(Protocol):
         man = self.client.get_ext_manifest(ext_handler, goal_state)
         return man.pkg_list
 
+    def get_in_vm_artifacts_profile(self):
+        logger.verbose("Get In-VM Artifacts Profile")
+        return self.client.get_in_vm_artifacts_profile()
+
     def report_provision_status(self, provision_status):
         validate_param("provision_status", provision_status, ProvisionStatus)
 
@@ -509,8 +513,8 @@ class WireClient(object):
         self.ext_conf = None
         self.last_request = 0
         self.req_count = 0
+        self.host_plugin = None
         self.status_blob = StatusBlob(self)
-        self.host_plugin = HostPluginProtocol(self.endpoint)
 
     def prevent_throttling(self):
         """
@@ -702,6 +706,9 @@ class WireClient(object):
                 self.update_shared_conf(goal_state)
                 self.update_certs(goal_state)
                 self.update_ext_conf(goal_state)
+                if self.host_plugin is not None:
+                    self.host_plugin.container_id = goal_state.container_id
+                    self.host_plugin.role_config_name = goal_state.role_instance_config_name
                 return
             except WireProtocolResourceGone:
                 logger.info("Incarnation is out of date. Update goalstate.")
@@ -795,8 +802,7 @@ class WireClient(object):
         ext_conf = self.get_ext_conf()
         if ext_conf.status_upload_blob is not None:
             if not self.status_blob.upload(ext_conf.status_upload_blob):
-                self.host_plugin.put_vm_status(self.status_blob,
-                                               ext_conf.status_upload_blob)
+                self.get_host_plugin().put_vm_status(self.status_blob, ext_conf.status_upload_blob)
 
     def report_role_prop(self, thumbprint):
         goal_state = self.get_goal_state()
@@ -911,6 +917,60 @@ class WireClient(object):
             "x-ms-guest-agent-public-x509-cert": cert
         }
 
+    def get_host_plugin(self):
+        if self.host_plugin is None:
+            goal_state = self.get_goal_state()
+            self.host_plugin = HostPluginProtocol(self.endpoint,
+                                                  goal_state.container_id,
+                                                  goal_state.role_instance_config_name)
+        return self.host_plugin
+
+    def get_in_vm_artifacts_profile(self):
+        ext_conf = self.ext_conf
+        if ext_conf and \
+                ext_conf.in_vm_artifacts_profile_blob and not \
+                ext_conf.in_vm_artifacts_profile_blob.isspace():
+            # try with the default protocol
+            in_vm_artifacts_profile_byte_arr = \
+                self._get_in_vm_artifacts_profile(ext_conf.in_vm_artifacts_profile_blob)
+
+            # try with host GA plugin
+            if in_vm_artifacts_profile_byte_arr is None:
+                logger.warn(
+                    "Failed to get extension artifacts profile using the default protocol. Try with HostGAPlugin")
+                host_plugin_endpoint, headers = \
+                    self.get_host_plugin().get_extension_artifact_url_and_headers(ext_conf.in_vm_artifacts_profile_blob)
+                in_vm_artifacts_profile_byte_arr = self._get_in_vm_artifacts_profile(host_plugin_endpoint, headers)
+
+            in_vm_artifacts_profile_json = self.decode_config(in_vm_artifacts_profile_byte_arr)
+            if in_vm_artifacts_profile_json and not in_vm_artifacts_profile_json.isspace():
+                return InVMArtifactsProfile(in_vm_artifacts_profile_json)
+
+    def _get_in_vm_artifacts_profile(self, blob_url, headers=None):
+        result = None
+        try:
+            resp = self.call_storage_service(restutil.http_get, blob_url, headers=headers)
+            response_body = resp.read()
+            if resp.status < 400:
+                decoded = self.decode_config(response_body)
+                if decoded and not decoded.isspace():
+                    result = InVMArtifactsProfile(decoded)
+            else:
+                logger.warn("Unexpected http status '{0}' is returned with the message '{1}'",
+                            resp.status,
+                            response_body)
+
+        except HttpError as e:
+            logger.warn("Encountered HttpError while getting InVMArtifactsProfile from '{0}' "
+                        "using the default protocol: [{1}]", blob_url, e)
+        else:
+            if resp.status == httpclient.OK:
+                result = resp.read()
+            else:
+                logger.warn("Failed to get InVMArtifactsProfile with status [{0}] from '{1}' "
+                            "using the default protocol", resp.status, blob_url)
+        return result
+
 
 class VersionInfo(object):
     def __init__(self, xml_text):
@@ -956,6 +1016,7 @@ class GoalState(object):
         self.certs_uri = None
         self.ext_uri = None
         self.role_instance_id = None
+        self.role_instance_config_name = None
         self.container_id = None
         self.load_balancer_probe_port = None
         self.parse(xml_text)
@@ -974,6 +1035,8 @@ class GoalState(object):
         self.ext_uri = findtext(xml_doc, "ExtensionsConfig")
         role_instance = find(xml_doc, "RoleInstance")
         self.role_instance_id = findtext(role_instance, "InstanceId")
+        role_config = find(role_instance, "Configuration")
+        self.role_instance_config_name = findtext(role_config, "ConfigName")
         container = find(xml_doc, "Container")
         self.container_id = findtext(container, "ContainerId")
         lbprobe_ports = find(xml_doc, "LBProbePorts")
@@ -1137,6 +1200,7 @@ class ExtensionsConfig(object):
         self.ext_handlers = ExtHandlerList()
         self.vmagent_manifests = VMAgentManifestList()
         self.status_upload_blob = None
+        self.in_vm_artifacts_profile_blob = None
         if xml_text is not None:
             self.parse(xml_text)
 
@@ -1171,6 +1235,7 @@ class ExtensionsConfig(object):
             self.parse_plugin_settings(ext_handler, plugin_settings)
 
         self.status_upload_blob = findtext(xml_doc, "StatusUploadBlob")
+        self.in_vm_artifacts_profile_blob = findtext(xml_doc, "InVMArtifactsProfileBlob")
 
     def parse_plugin(self, plugin):
         ext_handler = ExtHandler()
@@ -1272,3 +1337,27 @@ class ExtensionManifest(object):
 
             pkg.isinternal = isinternal
             self.pkg_list.versions.append(pkg)
+
+
+# Do not extend this class
+class InVMArtifactsProfile(object):
+    '''
+    deserialized json string of InVMArtifactsProfile.
+    It is expected to contain the following fields:
+    * inVMArtifactsProfileBlobSeqNo
+    * profileId (optional)
+    * onHold (optional)
+    * certificateThumbprint (optional)
+    * encryptedHealthChecks (optional)
+    * encryptedApplicationProfile (optional)
+    '''
+
+    def __init__(self, in_vm_artifacts_profile_json):
+        if in_vm_artifacts_profile_json and not in_vm_artifacts_profile_json.isspace():
+            self.__dict__.update(parse_json(in_vm_artifacts_profile_json))
+
+    def is_extension_handlers_handling_on_hold(self):
+        # hasattr() is not available in Python 2.6
+        if 'onHold' in self.__dict__:
+            return self.onHold.lower() == 'true'
+        return False
diff --git a/azurelinuxagent/common/utils/textutil.py b/azurelinuxagent/common/utils/textutil.py
index db320a62..f60df26a 100644
--- a/azurelinuxagent/common/utils/textutil.py
+++ b/azurelinuxagent/common/utils/textutil.py
@@ -278,15 +278,30 @@ def b64encode(s):
         return base64.b64encode(bytes(s, 'utf-8')).decode('utf-8')
     return base64.b64encode(s)
 
+
 def b64decode(s):
     from azurelinuxagent.common.version import PY_VERSION_MAJOR
     if PY_VERSION_MAJOR > 2:
         return base64.b64decode(s).decode('utf-8')
     return base64.b64decode(s)
 
+
 def safe_shlex_split(s):
     import shlex
     from azurelinuxagent.common.version import PY_VERSION
     if PY_VERSION[:2] == (2, 6):
         return shlex.split(s.encode('utf-8'))
     return shlex.split(s)
+
+
+def parse_json(json_str):
+    """
+    Parse json string and return a resulting dictionary
+    """
+    # trim null and whitespaces
+    result = None
+    if json_str and not json_str.isspace():
+        import json
+        result = json.loads(json_str.rstrip(' \t\r\n\0'))
+
+    return result
diff --git a/azurelinuxagent/ga/exthandlers.py b/azurelinuxagent/ga/exthandlers.py
index 192300e3..dda122c6 100644
--- a/azurelinuxagent/ga/exthandlers.py
+++ b/azurelinuxagent/ga/exthandlers.py
@@ -195,6 +195,12 @@ class ExtHandlersHandler(object):
             logger.verbose("No ext handler config found")
             return
 
+        in_vm_artifacts_profile = self.protocol.get_in_vm_artifacts_profile()
+        if in_vm_artifacts_profile and \
+           in_vm_artifacts_profile.is_extension_handlers_handling_on_hold():
+            logger.info("Handling Extension handlers is on hold")
+            return
+
         for ext_handler in self.ext_handlers.extHandlers:
             #TODO handle install in sequence, enable in parallel
             self.handle_ext_handler(ext_handler, etag)


Test patch:
diff --git a/tests/data/wire/goal_state.xml b/tests/data/wire/goal_state.xml
index 960444bb..0fd144f2 100644
--- a/tests/data/wire/goal_state.xml
+++ b/tests/data/wire/goal_state.xml
@@ -20,6 +20,7 @@
          <Certificates>http://certificatesuri/</Certificates>
          <ExtensionsConfig>http://extensionsconfiguri/</ExtensionsConfig>
          <FullConfig>http://fullconfiguri/</FullConfig>
+         <ConfigName>DummyRoleConfigName.xml</ConfigName>
          </Configuration>
        </RoleInstance>
      </RoleInstanceList>
diff --git a/tests/ga/test_extension.py b/tests/ga/test_extension.py
index fe30ca15..b7e50a30 100644
--- a/tests/ga/test_extension.py
+++ b/tests/ga/test_extension.py
@@ -286,7 +286,39 @@ class TestExtension(AgentTestCase):
         mock_fileutil.write_file.return_value = IOError("Mock IO Error")
         exthandlers_handler.run()
 
-    def _assert_ext_status(self, report_ext_status, expected_status, 
+    def test_handle_ext_handlers_on_hold_true(self, *args):
+        test_data = WireProtocolData(DATA_FILE)
+        exthandlers_handler, protocol = self._create_mock(test_data, *args)
+        exthandlers_handler.ext_handlers, exthandlers_handler.last_etag = protocol.get_ext_handlers()
+        protocol.get_in_vm_artifacts_profile = MagicMock()
+        exthandlers_handler.protocol = protocol
+
+        with patch.object(ExtHandlersHandler, 'handle_ext_handler') as patch_handle_ext_handler:
+            exthandlers_handler.handle_ext_handlers()
+            patch_handle_ext_handler.assert_not_called()
+
+    def test_handle_ext_handlers_on_hold_false(self, *args):
+        test_data = WireProtocolData(DATA_FILE)
+        exthandlers_handler, protocol = self._create_mock(test_data, *args)
+        exthandlers_handler.ext_handlers, exthandlers_handler.last_etag = protocol.get_ext_handlers()
+        exthandlers_handler.protocol = protocol
+
+        #Test when is_extension_handlers_handling_on_hold returns False
+        from azurelinuxagent.common.protocol.wire import InVMArtifactsProfile
+        mock_in_vm_artifacts_profile = InVMArtifactsProfile(MagicMock())
+        mock_in_vm_artifacts_profile.is_extension_handlers_handling_on_hold = Mock(return_value=False)
+        protocol.get_in_vm_artifacts_profile = Mock(return_value=mock_in_vm_artifacts_profile)
+        with patch.object(ExtHandlersHandler, 'handle_ext_handler') as patch_handle_ext_handler:
+            exthandlers_handler.handle_ext_handlers()
+            patch_handle_ext_handler.assert_called_once()
+
+        #Test when in_vm_artifacts_profile is not available
+        protocol.get_in_vm_artifacts_profile = Mock(return_value=None)
+        with patch.object(ExtHandlersHandler, 'handle_ext_handler') as patch_handle_ext_handler:
+            exthandlers_handler.handle_ext_handlers()
+            patch_handle_ext_handler.assert_called_once()
+
+    def _assert_ext_status(self, report_ext_status, expected_status,
                            expected_seq_no):
         self.assertTrue(report_ext_status.called)
         args, kw = report_ext_status.call_args
diff --git a/tests/protocol/test_hostplugin.py b/tests/protocol/test_hostplugin.py
index 296a1f58..ceabd2de 100644
--- a/tests/protocol/test_hostplugin.py
+++ b/tests/protocol/test_hostplugin.py
@@ -19,6 +19,8 @@ import unittest
 
 import azurelinuxagent.common.protocol.restapi as restapi
 import azurelinuxagent.common.protocol.wire as wire
+import azurelinuxagent.common.protocol.hostplugin as hostplugin
+from tests.protocol.mockwiredata import WireProtocolData, DATA_FILE
 from tests.tools import *
 
 wireserver_url = "168.63.129.16"
@@ -28,11 +30,16 @@ api_versions = '["2015-09-01"]'
 
 class TestHostPlugin(AgentTestCase):
     def test_fallback(self):
-        with patch.object(wire.HostPluginProtocol,
-                          "put_vm_status") as patch_put:
-            with patch.object(wire.StatusBlob, "upload") as patch_upload:
-                patch_upload.return_value = False
+        """
+        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)
+
+        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.upload_status_blob()
@@ -40,22 +47,26 @@ class TestHostPlugin(AgentTestCase):
                                 "Fallback was not engaged")
                 self.assertTrue(patch_put.call_args[0][1] == sas_url)
 
-    def test_fallback_data(self):
+    def test_validate_http_request(self):
+        """Validate correct set of data is sent to HostGAPlugin when reporting VM status"""
         from azurelinuxagent.common.protocol.hostplugin import API_VERSION
         from azurelinuxagent.common.utils import restutil
         exp_method = 'PUT'
         exp_url = 'http://{0}:32526/status'.format(wireserver_url)
-        exp_data = '{"content": "UmVhZHk=", "headers": [{"headerName": ' \
+        exp_data = '{"content": "eyJkdW1teSI6ICJkYXRhIn0=", "headers": [{"headerName": ' \
                    '"x-ms-version", "headerValue": "2014-02-14"}, ' \
                    '{"headerName": "x-ms-blob-type", "headerValue": null}], ' \
                    '"requestUri": "http://sas_url"}'
+        test_goal_state = wire.GoalState(WireProtocolData(DATA_FILE).goal_state)
+
         with patch.object(restutil, "http_request") as patch_http:
             wire_protocol_client = wire.WireProtocol(wireserver_url).client
-            plugin = wire_protocol_client.host_plugin
+            wire_protocol_client.get_goal_state = Mock(return_value=test_goal_state)
+            plugin = wire_protocol_client.get_host_plugin()
             blob = wire_protocol_client.status_blob
             blob.vm_status = restapi.VMStatus()
-            blob.vm_status.vmAgent.status = 'Ready'
-            with patch.object(plugin, "get_api_versions") as patch_api:
+            blob.data = '{"dummy": "data"}'
+            with patch.object(plugin, 'get_api_versions') as patch_api:
                 patch_api.return_value = API_VERSION
                 plugin.put_vm_status(blob, sas_url)
                 self.assertTrue(patch_http.call_count == 1)
@@ -63,7 +74,16 @@ class TestHostPlugin(AgentTestCase):
                 self.assertTrue(patch_http.call_args[0][1] == exp_url)
                 self.assertTrue(patch_http.call_args[0][2] == exp_data)
 
+                # Assert headers
+                headers = patch_http.call_args[1]['headers']
+                self.assertEqual(headers['x-ms-containerid'], test_goal_state.container_id)
+                self.assertEqual(headers['x-ms-host-config-name'], test_goal_state.role_instance_config_name)
+
     def test_no_fallback(self):
+        """
+        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:
             with patch.object(wire.StatusBlob, "upload") as patch_upload:
@@ -75,22 +95,29 @@ class TestHostPlugin(AgentTestCase):
                 self.assertTrue(patch_put.call_count == 0,
                                 "Fallback was engaged")
 
-    def test_init_put(self):
+    def test_validate_http_put(self):
+        """Validate correct set of data is sent to HostGAPlugin when reporting VM status"""
+        test_goal_state = wire.GoalState(WireProtocolData(DATA_FILE).goal_state)
         expected_url = "http://168.63.129.16:32526/status"
-        expected_headers = {'x-ms-version': '2015-09-01'}
-        expected_content = '{"content": "b2s=", ' \
+        expected_headers = {'x-ms-version': '2015-09-01',
+                            "Content-type": "application/json",
+                            "x-ms-containerid": test_goal_state.container_id,
+                            "x-ms-host-config-name": test_goal_state.role_instance_config_name}
+        expected_content = '{"content": "eyJkdW1teSI6ICJkYXRhIn0=", ' \
                            '"headers": [{"headerName": "x-ms-version", ' \
                            '"headerValue": "2014-02-14"}, ' \
                            '{"headerName": "x-ms-blob-type", "headerValue": ' \
                            '"BlockBlob"}], ' \
                            '"requestUri": "http://sas_url"}'
 
-        host_client = wire.HostPluginProtocol(wireserver_url)
+        host_client = wire.HostPluginProtocol(wireserver_url,
+                                              test_goal_state.container_id,
+                                              test_goal_state.role_instance_config_name)
         self.assertFalse(host_client.is_initialized)
         self.assertTrue(host_client.api_versions is None)
         status_blob = wire.StatusBlob(None)
         status_blob.vm_status = restapi.VMStatus()
-        status_blob.vm_status.vmAgent.status = 'ok'
+        status_blob.data = '{"dummy": "data"}'
         status_blob.type = "BlockBlob"
         with patch.object(wire.HostPluginProtocol,
                           "get_api_versions") as patch_get:
@@ -105,6 +132,36 @@ class TestHostPlugin(AgentTestCase):
                 self.assertTrue(patch_put.call_args[0][1] == expected_content)
                 self.assertTrue(patch_put.call_args[0][2] == expected_headers)
 
+    def test_validate_get_extension_artifacts(self):
+        test_goal_state = wire.GoalState(WireProtocolData(DATA_FILE).goal_state)
+        expected_url = hostplugin.URI_FORMAT_GET_EXTENSION_ARTIFACT.format(wireserver_url, hostplugin.HOST_PLUGIN_PORT)
+        expected_headers = {'x-ms-version': '2015-09-01',
+                            "x-ms-containerid": test_goal_state.container_id,
+                            "x-ms-host-config-name": test_goal_state.role_instance_config_name,
+                            "x-ms-artifact-location": sas_url}
+
+        host_client = wire.HostPluginProtocol(wireserver_url,
+                                              test_goal_state.container_id,
+                                              test_goal_state.role_instance_config_name)
+        self.assertFalse(host_client.is_initialized)
+        self.assertTrue(host_client.api_versions is None)
+
+        with patch.object(wire.HostPluginProtocol, "get_api_versions", return_value=api_versions) as patch_get:
+            actual_url, actual_headers = host_client.get_extension_artifact_url_and_headers(sas_url)
+            self.assertTrue(host_client.is_initialized)
+            self.assertFalse(host_client.api_versions is None)
+            self.assertEqual(expected_url, actual_url)
+            for k in expected_headers:
+                self.assertTrue(k in actual_headers)
+                self.assertEqual(expected_headers[k], actual_headers[k])
+
+class MockResponse:
+    def __init__(self, body, status_code):
+        self.body = body
+        self.status = status_code
+
+    def read(self):
+        return self.body
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/protocol/test_wire.py b/tests/protocol/test_wire.py
index 59af98b9..bc2f4964 100644
--- a/tests/protocol/test_wire.py
+++ b/tests/protocol/test_wire.py
@@ -15,11 +15,12 @@
 # Requires Python 2.4+ and Openssl 1.0+
 #
 
-from azurelinuxagent.common.protocol.wire import WireProtocol, WireClient
+from azurelinuxagent.common.protocol.wire import *
 from tests.protocol.mockwiredata import *
 
 data_with_bom = b'\xef\xbb\xbfhehe'
-
+testurl = 'http://foo'
+wireserver_url = '168.63.129.16'
 
 @patch("time.sleep")
 @patch("azurelinuxagent.common.protocol.wire.CryptUtil")
@@ -29,7 +30,7 @@ class TestWireProtocolGetters(AgentTestCase):
         mock_restutil.http_get.side_effect = test_data.mock_http_get
         MockCryptUtil.side_effect = test_data.mock_crypt_util
 
-        protocol = WireProtocol("foo.bar")
+        protocol = WireProtocol(wireserver_url)
         protocol.detect()
         protocol.get_vminfo()
         protocol.get_certs()
@@ -75,7 +76,7 @@ class TestWireProtocolGetters(AgentTestCase):
         from azurelinuxagent.common.utils import restutil
         with patch.object(restutil, 'http_get') as http_patch:
             http_req = restutil.http_get
-            url = 'http://foo'
+            url = testurl
             headers = {}
 
             # no kwargs
@@ -99,5 +100,121 @@ class TestWireProtocolGetters(AgentTestCase):
             for c in http_patch.call_args_list:
                 self.assertTrue(c[-1]['chk_proxy'] == True)
 
+    def test_get_host_ga_plugin(self, *args):
+        wire_protocol_client = WireProtocol(wireserver_url).client
+        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+
+        with patch.object(WireClient, "get_goal_state", return_value = goal_state) as patch_get_goal_state:
+            host_plugin = wire_protocol_client.get_host_plugin()
+            self.assertEqual(goal_state.container_id, host_plugin.container_id)
+            self.assertEqual(goal_state.role_instance_config_name, host_plugin.role_config_name)
+            patch_get_goal_state.assert_called_once()
+
+    def test_upload_status_blob_default(self, *args):
+        wire_protocol_client = WireProtocol(wireserver_url).client
+        wire_protocol_client.ext_conf = ExtensionsConfig(None)
+        wire_protocol_client.ext_conf.status_upload_blob = testurl
+
+        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:
+                    wire_protocol_client.upload_status_blob()
+
+                    patch_default_upload.assert_called_once_with(testurl)
+                    patch_get_goal_state.assert_not_called()
+                    patch_host_ga_plugin_upload.assert_not_called()
+
+    def test_upload_status_blob_host_ga_plugin(self, *args):
+        wire_protocol_client = WireProtocol(wireserver_url).client
+        wire_protocol_client.ext_conf = ExtensionsConfig(None)
+        wire_protocol_client.ext_conf.status_upload_blob = testurl
+        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.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)
+
+    def test_get_in_vm_artifacts_profile_blob_not_available(self, *args):
+        wire_protocol_client = WireProtocol(wireserver_url).client
+        wire_protocol_client.ext_conf = ExtensionsConfig(None)
+
+        # Test when in_vm_artifacts_profile_blob is null/None
+        self.assertEqual(None, wire_protocol_client.get_in_vm_artifacts_profile())
+
+        #Test when in_vm_artifacts_profile_blob is whitespace
+        wire_protocol_client.ext_conf.in_vm_artifacts_profile_blob = "  "
+        self.assertEqual(None, wire_protocol_client.get_in_vm_artifacts_profile())
+
+    def test_get_in_vm_artifacts_profile_response_body_not_valid(self, *args):
+        wire_protocol_client = WireProtocol(wireserver_url).client
+        wire_protocol_client.ext_conf = ExtensionsConfig(None)
+        wire_protocol_client.ext_conf.in_vm_artifacts_profile_blob = testurl
+        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+        wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
+
+        with patch.object(HostPluginProtocol, "get_extension_artifact_url_and_headers",
+                          return_value = ['dummy_url', {}]) as host_plugin_get_artifact_url_and_headers:
+            #Test when response body is None
+            wire_protocol_client.call_storage_service = Mock(return_value=MockResponse(None, 200))
+            in_vm_artifacts_profile = wire_protocol_client.get_in_vm_artifacts_profile()
+            self.assertTrue(in_vm_artifacts_profile is None)
+
+            #Test when response body is None
+            wire_protocol_client.call_storage_service = Mock(return_value=MockResponse('   '.encode('utf-8'), 200))
+            in_vm_artifacts_profile = wire_protocol_client.get_in_vm_artifacts_profile()
+            self.assertTrue(in_vm_artifacts_profile is None)
+
+            #Test when response body is None
+            wire_protocol_client.call_storage_service = Mock(return_value=MockResponse('{ }'.encode('utf-8'), 200))
+            in_vm_artifacts_profile = wire_protocol_client.get_in_vm_artifacts_profile()
+            self.assertEqual(dict(), in_vm_artifacts_profile.__dict__,
+                             'If in_vm_artifacts_profile_blob has empty json dictionary, in_vm_artifacts_profile '
+                             'should contain nothing')
+
+            host_plugin_get_artifact_url_and_headers.assert_called_with(testurl)
+
+
+    def test_get_in_vm_artifacts_profile_default(self, *args):
+        wire_protocol_client = WireProtocol(wireserver_url).client
+        wire_protocol_client.ext_conf = ExtensionsConfig(None)
+        wire_protocol_client.ext_conf.in_vm_artifacts_profile_blob = testurl
+        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+        wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
+
+        wire_protocol_client.call_storage_service = Mock(return_value=MockResponse('{"onHold": "true"}'.encode('utf-8'), 200))
+        in_vm_artifacts_profile = wire_protocol_client.get_in_vm_artifacts_profile()
+        self.assertEqual(dict(onHold='true'), in_vm_artifacts_profile.__dict__)
+        self.assertTrue(in_vm_artifacts_profile.is_extension_handlers_handling_on_hold())
+
+    def test_get_in_vm_artifacts_profile_host_ga_plugin(self, *args):
+        wire_protocol_client = WireProtocol(wireserver_url).client
+        wire_protocol_client.ext_conf = ExtensionsConfig(None)
+        wire_protocol_client.ext_conf.in_vm_artifacts_profile_blob = testurl
+        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+        wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
+
+        wire_protocol_client._get_in_vm_artifacts_profile = Mock(side_effect=[None, '{"onHold": "true"}'.encode('utf-8')])
+
+        with patch.object(HostPluginProtocol, "get_extension_artifact_url_and_headers",
+                          return_value = ['dummy_url', {}]) as host_plugin_get_artifact_url_and_headers:
+            in_vm_artifacts_profile = wire_protocol_client.get_in_vm_artifacts_profile()
+            self.assertEqual(dict(onHold='true'), in_vm_artifacts_profile.__dict__)
+            self.assertTrue(in_vm_artifacts_profile.is_extension_handlers_handling_on_hold())
+            host_plugin_get_artifact_url_and_headers.assert_called_once_with(testurl)
+
+
+class MockResponse:
+    def __init__(self, body, status_code):
+        self.body = body
+        self.status = status_code
+
+    def read(self):
+        return self.body
+
 if __name__ == '__main__':
     unittest.main()