Repository: Azure/WALinuxAgent
Pull request #1777 (2020-02-11 20:38:47)

Separate goal state from protocol
if an error occurs while _update_from_goal_state is populating the goal state the retry loop may end up producing a goal state that is only partially populated.

This PR addresses that issue and starts to refactor the goal state out of the protocol. I will do further refactoring to improve that separation once we remove the metadata protocol.

The code refactored in this PR is heavily exercised by the current tests. I will add tests specific to the goal state class in my next round of refactoring.

I did manual verification using the debugger on a live VM, and automation reported no errors.

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/azure/walinuxagent/1777)
<!-- Reviewable:end -->


codecov[bot]: # [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1777?src=pr&el=h1) Report
> :exclamation: No coverage uploaded for pull request base (`develop@d55c1a9`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference#section-missing-base-commit).
> The diff coverage is `92.75%`.

[![Impacted file tree graph](https://codecov.io/gh/Azure/WALinuxAgent/pull/1777/graphs/tree.svg?width=650&token=ae90E2KYCw&height=150&src=pr)](https://codecov.io/gh/Azure/WALinuxAgent/pull/1777?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##             develop    #1777   +/-   ##
==========================================
  Coverage           ?   68.63%           
==========================================
  Files              ?       82           
  Lines              ?    11685           
  Branches           ?     1640           
==========================================
  Hits               ?     8020           
  Misses             ?     3319           
  Partials           ?      346
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/1777?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/protocol/goal\_state.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1777/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC9nb2FsX3N0YXRlLnB5) | `95.01% <100%> (ø)` | |
| [azurelinuxagent/common/protocol/wire.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1777/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC93aXJlLnB5) | `71.79% <78.57%> (ø)` | |
| [azurelinuxagent/ga/monitor.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1777/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2dhL21vbml0b3IucHk=) | `88.95% <96.15%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1777?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1777?src=pr&el=footer). Last update [d55c1a9...01a11fb](https://codecov.io/gh/Azure/WALinuxAgent/pull/1777?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/azurelinuxagent/common/protocol/__init__.py b/azurelinuxagent/common/protocol/__init__.py
index 746fa076..24121bff 100644
--- a/azurelinuxagent/common/protocol/__init__.py
+++ b/azurelinuxagent/common/protocol/__init__.py
@@ -14,8 +14,3 @@
 #
 # Requires Python 2.6+ and Openssl 1.0+
 #
-
-from azurelinuxagent.common.protocol.util import get_protocol_util, \
-                                                 OVF_FILE_NAME
-
-
diff --git a/azurelinuxagent/common/protocol/goal_state.py b/azurelinuxagent/common/protocol/goal_state.py
new file mode 100644
index 00000000..68cf1b5b
--- /dev/null
+++ b/azurelinuxagent/common/protocol/goal_state.py
@@ -0,0 +1,411 @@
+# Microsoft Azure Linux Agent
+#
+# Copyright 2020 Microsoft Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Requires Python 2.6+ and Openssl 1.0+
+
+import json
+import os
+import re
+
+import azurelinuxagent.common.conf as conf
+import azurelinuxagent.common.logger as logger
+from azurelinuxagent.common.datacontract import set_properties, DataContract, DataContractList
+from azurelinuxagent.common.event import CONTAINER_ID_ENV_VARIABLE
+from azurelinuxagent.common.utils import fileutil
+from azurelinuxagent.common.utils.cryptutil import CryptUtil
+from azurelinuxagent.common.utils.textutil import parse_doc, findall, find, findtext, getattrib, gettext
+from azurelinuxagent.common.protocol.restapi import *
+
+GOAL_STATE_URI = "http://{0}/machine/?comp=goalstate"
+CERTS_FILE_NAME = "Certificates.xml"
+P7M_FILE_NAME = "Certificates.p7m"
+PEM_FILE_NAME = "Certificates.pem"
+TRANSPORT_CERT_FILE_NAME = "TransportCert.pem"
+TRANSPORT_PRV_FILE_NAME = "TransportPrivate.pem"
+
+# Store the last retrieved container id as an environment variable to be shared between threads for telemetry purposes
+CONTAINER_ID_ENV_VARIABLE = "AZURE_GUEST_AGENT_CONTAINER_ID"
+
+
+class GoalState(object):
+    def __init__(self, wire_client, full_goal_state=False, base_incarnation=None):
+        """
+        Fetches the goal state using the given wire client.
+
+        By default it fetches only the goal state itself; to fetch the entire goal state (that includes all the
+        nested components, such as the extension config) use the 'full_goal_state' parameter.
+
+        If 'base_incarnation' is given, it fetches the full goal state if the new incarnation is different than
+        the given value, otherwise it fetches only the goal state itself.
+
+        For better code readability, use the static fetch_* methods below instead of instantiating GoalState
+        directly.
+
+        """
+        uri = GOAL_STATE_URI.format(wire_client.get_endpoint())
+        self.xml_text = wire_client.fetch_config(uri, wire_client.get_header())
+        xml_doc = parse_doc(self.xml_text)
+
+        self.incarnation = findtext(xml_doc, "Incarnation")
+        self.expected_state = findtext(xml_doc, "ExpectedState")
+        role_instance = find(xml_doc, "RoleInstance")
+        self.role_instance_id = findtext(role_instance, "InstanceId")
+        role_config = find(role_instance, "Configuration")
+        self.role_config_name = findtext(role_config, "ConfigName")
+        container = find(xml_doc, "Container")
+        self.container_id = findtext(container, "ContainerId")
+        lbprobe_ports = find(xml_doc, "LBProbePorts")
+        self.load_balancer_probe_port = findtext(lbprobe_ports, "Port")
+
+        os.environ[CONTAINER_ID_ENV_VARIABLE] = self.container_id
+
+        if not (full_goal_state or base_incarnation is not None and self.incarnation != base_incarnation):
+            self.hosting_env = None
+            self.shared_conf = None
+            self.certs = None
+            self.ext_conf = None
+            self.remote_access = None
+            return
+
+        uri = findtext(xml_doc, "HostingEnvironmentConfig")
+        xml_text = wire_client.fetch_config(uri, wire_client.get_header())
+        self.hosting_env = HostingEnv(xml_text)
+
+        uri = findtext(xml_doc, "SharedConfig")
+        xml_text = wire_client.fetch_config(uri, wire_client.get_header())
+        self.shared_conf = SharedConfig(xml_text)
+
+        uri = findtext(xml_doc, "Certificates")
+        if uri is None:
+            self.certs = None
+        else:
+            xml_text = wire_client.fetch_config(uri, wire_client.get_header_for_cert())
+            self.certs = Certificates(xml_text)
+
+        uri = findtext(xml_doc, "ExtensionsConfig")
+        if uri is None:
+            self.ext_conf = ExtensionsConfig(None)
+        else:
+            xml_text = wire_client.fetch_config(uri, wire_client.get_header())
+            self.ext_conf = ExtensionsConfig(xml_text)
+
+        uri = findtext(container, "RemoteAccessInfo")
+        if uri is None:
+            self.remote_access = None
+        else:
+            xml_text = wire_client.fetch_config(uri, wire_client.get_header_for_cert())
+            self.remote_access = RemoteAccess(xml_text)
+
+    @staticmethod
+    def fetch_goal_state(wire_client):
+        """
+        Fetches the goal state, not including any nested properties (such as extension config).
+        """
+        return GoalState(wire_client)
+
+    @staticmethod
+    def fetch_full_goal_state(wire_client):
+        """
+        Fetches the full goal state, including nested properties (such as extension config).
+        """
+        return GoalState(wire_client, full_goal_state=True)
+
+    @staticmethod
+    def fetch_full_goal_state_if_incarnation_different_than(wire_client, incarnation):
+        """
+        Fetches the full goal state if the new incarnation is different than 'incarnation', otherwise returns None.
+        """
+        goal_state = GoalState(wire_client, base_incarnation=incarnation)
+        return goal_state if goal_state.incarnation != incarnation else None
+
+
+class HostingEnv(object):
+    def __init__(self, xml_text):
+        self.xml_text = xml_text
+        xml_doc = parse_doc(xml_text)
+        incarnation = find(xml_doc, "Incarnation")
+        self.vm_name = getattrib(incarnation, "instance")
+        role = find(xml_doc, "Role")
+        self.role_name = getattrib(role, "name")
+        deployment = find(xml_doc, "Deployment")
+        self.deployment_name = getattrib(deployment, "name")
+
+
+class SharedConfig(object):
+    def __init__(self, xml_text):
+        self.xml_text = xml_text
+
+
+class Certificates(object):
+    def __init__(self, xml_text):
+        self.cert_list = CertList()
+
+        # Save the certificates
+        local_file = os.path.join(conf.get_lib_dir(), CERTS_FILE_NAME)
+        fileutil.write_file(local_file, xml_text)
+
+        # Separate the certificates into individual files.
+        xml_doc = parse_doc(xml_text)
+        data = findtext(xml_doc, "Data")
+        if data is None:
+            return
+
+        # if the certificates format is not Pkcs7BlobWithPfxContents do not parse it
+        certificateFormat = findtext(xml_doc, "Format")
+        if certificateFormat and certificateFormat != "Pkcs7BlobWithPfxContents":
+            logger.warn("The Format is not Pkcs7BlobWithPfxContents. Format is " + certificateFormat)
+            return
+
+        cryptutil = CryptUtil(conf.get_openssl_cmd())
+        p7m_file = os.path.join(conf.get_lib_dir(), P7M_FILE_NAME)
+        p7m = ("MIME-Version:1.0\n"
+               "Content-Disposition: attachment; filename=\"{0}\"\n"
+               "Content-Type: application/x-pkcs7-mime; name=\"{1}\"\n"
+               "Content-Transfer-Encoding: base64\n"
+               "\n"
+               "{2}").format(p7m_file, p7m_file, data)
+
+        fileutil.write_file(p7m_file, p7m)
+
+        trans_prv_file = os.path.join(conf.get_lib_dir(), TRANSPORT_PRV_FILE_NAME)
+        trans_cert_file = os.path.join(conf.get_lib_dir(), TRANSPORT_CERT_FILE_NAME)
+        pem_file = os.path.join(conf.get_lib_dir(), PEM_FILE_NAME)
+        # decrypt certificates
+        cryptutil.decrypt_p7m(p7m_file, trans_prv_file, trans_cert_file, pem_file)
+
+        # The parsing process use public key to match prv and crt.
+        buf = []
+        begin_crt = False
+        begin_prv = False
+        prvs = {}
+        thumbprints = {}
+        index = 0
+        v1_cert_list = []
+        with open(pem_file) as pem:
+            for line in pem.readlines():
+                buf.append(line)
+                if re.match(r'[-]+BEGIN.*KEY[-]+', line):
+                    begin_prv = True
+                elif re.match(r'[-]+BEGIN.*CERTIFICATE[-]+', line):
+                    begin_crt = True
+                elif re.match(r'[-]+END.*KEY[-]+', line):
+                    tmp_file = Certificates._write_to_tmp_file(index, 'prv', buf)
+                    pub = cryptutil.get_pubkey_from_prv(tmp_file)
+                    prvs[pub] = tmp_file
+                    buf = []
+                    index += 1
+                    begin_prv = False
+                elif re.match(r'[-]+END.*CERTIFICATE[-]+', line):
+                    tmp_file = Certificates._write_to_tmp_file(index, 'crt', buf)
+                    pub = cryptutil.get_pubkey_from_crt(tmp_file)
+                    thumbprint = cryptutil.get_thumbprint_from_crt(tmp_file)
+                    thumbprints[pub] = thumbprint
+                    # Rename crt with thumbprint as the file name
+                    crt = "{0}.crt".format(thumbprint)
+                    v1_cert_list.append({
+                        "name": None,
+                        "thumbprint": thumbprint
+                    })
+                    os.rename(tmp_file, os.path.join(conf.get_lib_dir(), crt))
+                    buf = []
+                    index += 1
+                    begin_crt = False
+
+        # Rename prv key with thumbprint as the file name
+        for pubkey in prvs:
+            thumbprint = thumbprints[pubkey]
+            if thumbprint:
+                tmp_file = prvs[pubkey]
+                prv = "{0}.prv".format(thumbprint)
+                os.rename(tmp_file, os.path.join(conf.get_lib_dir(), prv))
+                logger.info("Found private key matching thumbprint {0}".format(thumbprint))
+            else:
+                # Since private key has *no* matching certificate,
+                # it will not be named correctly
+                logger.warn("Found NO matching cert/thumbprint for private key!")
+
+        # Log if any certificates were found without matching private keys
+        # This can happen (rarely), and is useful to know for debugging
+        for pubkey in thumbprints:
+            if not pubkey in prvs:
+                msg = "Certificate with thumbprint {0} has no matching private key."
+                logger.info(msg.format(thumbprints[pubkey]))
+
+        for v1_cert in v1_cert_list:
+            cert = Cert()
+            set_properties("certs", cert, v1_cert)
+            self.cert_list.certificates.append(cert)
+
+    @staticmethod
+    def _write_to_tmp_file(index, suffix, buf):
+        file_name = os.path.join(conf.get_lib_dir(), "{0}.{1}".format(index, suffix))
+        fileutil.write_file(file_name, "".join(buf))
+        return file_name
+
+
+class ExtensionsConfig(object):
+    def __init__(self, xml_text):
+        self.xml_text = xml_text
+        self.ext_handlers = ExtHandlerList()
+        self.vmagent_manifests = VMAgentManifestList()
+        self.status_upload_blob = None
+        self.status_upload_blob_type = None
+        self.artifacts_profile_blob = None
+
+        if xml_text is None:
+            return
+
+        xml_doc = parse_doc(self.xml_text)
+
+        ga_families_list = find(xml_doc, "GAFamilies")
+        ga_families = findall(ga_families_list, "GAFamily")
+
+        for ga_family in ga_families:
+            family = findtext(ga_family, "Name")
+            uris_list = find(ga_family, "Uris")
+            uris = findall(uris_list, "Uri")
+            manifest = VMAgentManifest()
+            manifest.family = family
+            for uri in uris:
+                manifestUri = VMAgentManifestUri(uri=gettext(uri))
+                manifest.versionsManifestUris.append(manifestUri)
+            self.vmagent_manifests.vmAgentManifests.append(manifest)
+
+        plugins_list = find(xml_doc, "Plugins")
+        plugins = findall(plugins_list, "Plugin")
+        plugin_settings_list = find(xml_doc, "PluginSettings")
+        plugin_settings = findall(plugin_settings_list, "Plugin")
+
+        for plugin in plugins:
+            ext_handler = ExtensionsConfig._parse_plugin(plugin)
+            self.ext_handlers.extHandlers.append(ext_handler)
+            ExtensionsConfig._parse_plugin_settings(ext_handler, plugin_settings)
+
+        self.status_upload_blob = findtext(xml_doc, "StatusUploadBlob")
+        self.artifacts_profile_blob = findtext(xml_doc, "InVMArtifactsProfileBlob")
+
+        status_upload_node = find(xml_doc, "StatusUploadBlob")
+        self.status_upload_blob_type = getattrib(status_upload_node, "statusBlobType")
+        logger.verbose("Extension config shows status blob type as [{0}]", self.status_upload_blob_type)
+
+    @staticmethod
+    def _parse_plugin(plugin):
+        ext_handler = ExtHandler()
+        ext_handler.name = getattrib(plugin, "name")
+        ext_handler.properties.version = getattrib(plugin, "version")
+        ext_handler.properties.state = getattrib(plugin, "state")
+
+        location = getattrib(plugin, "location")
+        failover_location = getattrib(plugin, "failoverlocation")
+        for uri in [location, failover_location]:
+            version_uri = ExtHandlerVersionUri()
+            version_uri.uri = uri
+            ext_handler.versionUris.append(version_uri)
+        return ext_handler
+
+    @staticmethod
+    def _parse_plugin_settings(ext_handler, plugin_settings):
+        if plugin_settings is None:
+            return
+
+        name = ext_handler.name
+        version = ext_handler.properties.version
+        settings = [x for x in plugin_settings \
+                    if getattrib(x, "name") == name and \
+                    getattrib(x, "version") == version]
+
+        if settings is None or len(settings) == 0:
+            return
+
+        runtime_settings = None
+        runtime_settings_node = find(settings[0], "RuntimeSettings")
+        seqNo = getattrib(runtime_settings_node, "seqNo")
+        runtime_settings_str = gettext(runtime_settings_node)
+        try:
+            runtime_settings = json.loads(runtime_settings_str)
+        except ValueError as e:
+            logger.error("Invalid extension settings")
+            return
+
+        depends_on_level = 0
+        depends_on_node = find(settings[0], "DependsOn")
+        if depends_on_node != None:
+            try:
+                depends_on_level = int(getattrib(depends_on_node, "dependencyLevel"))
+            except (ValueError, TypeError):
+                logger.warn("Could not parse dependencyLevel for handler {0}. Setting it to 0".format(name))
+                depends_on_level = 0
+
+        for plugin_settings_list in runtime_settings["runtimeSettings"]:
+            handler_settings = plugin_settings_list["handlerSettings"]
+            ext = Extension()
+            # There is no "extension name" in wire protocol.
+            # Put
+            ext.name = ext_handler.name
+            ext.sequenceNumber = seqNo
+            ext.publicSettings = handler_settings.get("publicSettings")
+            ext.protectedSettings = handler_settings.get("protectedSettings")
+            ext.dependencyLevel = depends_on_level
+            thumbprint = handler_settings.get(
+                "protectedSettingsCertThumbprint")
+            ext.certificateThumbprint = thumbprint
+            ext_handler.properties.extensions.append(ext)
+
+
+class RemoteAccess(object):
+    """
+    Object containing information about user accounts
+    """
+    #
+    # <RemoteAccess>
+    #   <Version/>
+    #   <Incarnation/>
+    #    <Users>
+    #       <User>
+    #         <Name/>
+    #         <Password/>
+    #         <Expiration/>
+    #       </User>
+    #     </Users>
+    #   </RemoteAccess>
+    #
+    def __init__(self, xml_text):
+        self.xml_text = xml_text
+        self.version = None
+        self.incarnation = None
+        self.user_list = RemoteAccessUsersList()
+
+        if self.xml_text is None or len(self.xml_text) == 0:
+            return
+
+        xml_doc = parse_doc(self.xml_text)
+        self.version = findtext(xml_doc, "Version")
+        self.incarnation = findtext(xml_doc, "Incarnation")
+        user_collection = find(xml_doc, "Users")
+        users = findall(user_collection, "User")
+
+        for user in users:
+            remote_access_user = RemoteAccess._parse_user(user)
+            self.user_list.users.append(remote_access_user)
+
+    @staticmethod
+    def _parse_user(user):
+        name = findtext(user, "Name")
+        encrypted_password = findtext(user, "Password")
+        expiration = findtext(user, "Expiration")
+        remote_access_user = RemoteAccessUser(name, encrypted_password, expiration)
+        return remote_access_user
+
diff --git a/azurelinuxagent/common/protocol/wire.py b/azurelinuxagent/common/protocol/wire.py
index 43f29c7f..1cfdd6fa 100644
--- a/azurelinuxagent/common/protocol/wire.py
+++ b/azurelinuxagent/common/protocol/wire.py
@@ -20,20 +20,19 @@ import datetime
 import json
 import os
 import random
-import re
 import time
 import xml.sax.saxutils as saxutils
 from datetime import datetime
-from threading import Lock
 
 import azurelinuxagent.common.conf as conf
 import azurelinuxagent.common.logger as logger
 import azurelinuxagent.common.utils.textutil as textutil
-from azurelinuxagent.common.datacontract import validate_param, set_properties
-from azurelinuxagent.common.event import add_periodic, WALAEventOperation, CONTAINER_ID_ENV_VARIABLE
+from azurelinuxagent.common.datacontract import validate_param
+from azurelinuxagent.common.event import add_periodic, WALAEventOperation
 from azurelinuxagent.common.exception import ProtocolNotFoundError, \
     ResourceGoneError, ExtensionDownloadError, InvalidContainerError, ProtocolError, HttpError
 from azurelinuxagent.common.future import httpclient, bytebuffer
+from azurelinuxagent.common.protocol.goal_state import GoalState, TRANSPORT_CERT_FILE_NAME, TRANSPORT_PRV_FILE_NAME
 from azurelinuxagent.common.protocol.hostplugin import HostPluginProtocol
 from azurelinuxagent.common.protocol.restapi import *
 from azurelinuxagent.common.telemetryevent import TelemetryEventList
@@ -41,11 +40,10 @@ from azurelinuxagent.common.utils import fileutil, restutil
 from azurelinuxagent.common.utils.archive import StateFlusher
 from azurelinuxagent.common.utils.cryptutil import CryptUtil
 from azurelinuxagent.common.utils.textutil import parse_doc, findall, find, \
-    findtext, getattrib, gettext, remove_bom, get_bytes_from_pem, parse_json
+    findtext, gettext, remove_bom, get_bytes_from_pem, parse_json
 from azurelinuxagent.common.version import AGENT_NAME, CURRENT_VERSION
 
 VERSION_INFO_URI = "http://{0}/?comp=versions"
-GOAL_STATE_URI = "http://{0}/machine/?comp=goalstate"
 HEALTH_REPORT_URI = "http://{0}/machine?comp=health"
 ROLE_PROP_URI = "http://{0}/machine?comp=roleProperties"
 TELEMETRY_URI = "http://{0}/machine?comp=telemetrydata"
@@ -55,16 +53,9 @@ INCARNATION_FILE_NAME = "Incarnation"
 GOAL_STATE_FILE_NAME = "GoalState.{0}.xml"
 HOSTING_ENV_FILE_NAME = "HostingEnvironmentConfig.xml"
 SHARED_CONF_FILE_NAME = "SharedConfig.xml"
-CERTS_FILE_NAME = "Certificates.xml"
 REMOTE_ACCESS_FILE_NAME = "RemoteAccess.{0}.xml"
-P7M_FILE_NAME = "Certificates.p7m"
-PEM_FILE_NAME = "Certificates.pem"
 EXT_CONF_FILE_NAME = "ExtensionsConfig.{0}.xml"
 MANIFEST_FILE_NAME = "{0}.{1}.manifest.xml"
-TRANSPORT_CERT_FILE_NAME = "TransportCert.pem"
-TRANSPORT_PRV_FILE_NAME = "TransportPrivate.pem"
-# Store the last retrieved container id as an environment variable to be shared between threads for telemetry purposes
-CONTAINER_ID_ENV_VARIABLE = "AZURE_GUEST_AGENT_CONTAINER_ID"
 
 PROTOCOL_VERSION = "2012-11-30"
 ENDPOINT_FINE_NAME = "WireServer"
@@ -147,8 +138,7 @@ class WireProtocol(Protocol):
 
     def get_ext_handler_pkgs(self, ext_handler):
         logger.verbose("Get extension handler package")
-        goal_state = self.client.get_goal_state()
-        man = self.client.get_ext_manifest(ext_handler, goal_state)
+        man = self.client.get_ext_manifest(ext_handler)
         return man.pkg_list
 
     def get_artifacts_profile(self):
@@ -694,60 +684,6 @@ class WireClient(object):
                 raise
         return resp
 
-    def update_hosting_env(self, goal_state):
-        if goal_state.hosting_env_uri is None:
-            raise ProtocolError("HostingEnvironmentConfig uri is empty")
-        local_file = os.path.join(conf.get_lib_dir(), HOSTING_ENV_FILE_NAME)
-        xml_text = self.fetch_config(goal_state.hosting_env_uri,
-                                     self.get_header())
-        self.save_cache(local_file, xml_text)
-        self._set_hosting_env(HostingEnv(xml_text))
-
-    def update_shared_conf(self, goal_state):
-        if goal_state.shared_conf_uri is None:
-            raise ProtocolError("SharedConfig uri is empty")
-        local_file = os.path.join(conf.get_lib_dir(), SHARED_CONF_FILE_NAME)
-        xml_text = self.fetch_config(goal_state.shared_conf_uri,
-                                     self.get_header())
-        self.save_cache(local_file, xml_text)
-        self._set_shared_conf(SharedConfig(xml_text))
-
-    def update_certs(self, goal_state):
-        if goal_state.certs_uri is None:
-            return
-        local_file = os.path.join(conf.get_lib_dir(), CERTS_FILE_NAME)
-        xml_text = self.fetch_config(goal_state.certs_uri,
-                                     self.get_header_for_cert())
-        self.save_cache(local_file, xml_text)
-        self._set_certs(Certificates(self, xml_text))
-
-    def update_remote_access_conf(self, goal_state):
-        if goal_state.remote_access_uri is None:
-            # Nothing in accounts data.  Just return, nothing to do.
-            return
-        xml_text = self.fetch_config(goal_state.remote_access_uri,
-                                     self.get_header_for_cert())
-        remote_access = RemoteAccess(xml_text)
-        local_file = os.path.join(conf.get_lib_dir(), REMOTE_ACCESS_FILE_NAME.format(remote_access.incarnation))
-        self.save_cache(local_file, xml_text)
-        self._set_remote_access(remote_access)
-
-    def get_remote_access(self):
-        # This property can be None in the GoalState, so not raising a ProtocolError here
-        return self._remote_access
-
-    def update_ext_conf(self, goal_state):
-        if goal_state.ext_uri is None:
-            logger.info("ExtensionsConfig.xml uri is empty")
-            self._set_ext_conf(ExtensionsConfig(None))
-            return
-        incarnation = goal_state.incarnation
-        local_file = os.path.join(conf.get_lib_dir(),
-                                  EXT_CONF_FILE_NAME.format(incarnation))
-        xml_text = self.fetch_config(goal_state.ext_uri, self.get_header())
-        self.save_cache(local_file, xml_text)
-        self._set_ext_conf(ExtensionsConfig(xml_text))
-
     # Type of update performed by _update_from_goal_state()
     class _UpdateType(object):
         # Update the Host GA Plugin client (Container ID and RoleConfigName)
@@ -774,42 +710,25 @@ class WireClient(object):
         """
         Fetches a new goal state and updates the internal state of the WireClient according to the requested 'refresh_type'
         """
-        uri = GOAL_STATE_URI.format(self.get_endpoint())
-
         max_retry = 3
 
         for retry in range(1, max_retry + 1):
             try:
-                new_goal_state_xml = self.fetch_config(uri, self.get_header())
-                new_goal_state = GoalState(new_goal_state_xml)
-
-                def update_host_plugin():
-                    if self._host_plugin is not None:
-                        self._host_plugin.update_container_id(new_goal_state.container_id)
-                        self._host_plugin.update_role_config_name(new_goal_state.role_config_name)
-
                 if refresh_type == WireClient._UpdateType.HostPlugin:
-                    update_host_plugin()
+                    goal_state = GoalState.fetch_goal_state(self)
+                    self._update_host_plugin(goal_state.container_id, goal_state.role_config_name)
                     return
 
-                if self._goal_state is None or refresh_type == WireClient._UpdateType.GoalStateForced or \
-                        new_goal_state.incarnation != self._goal_state.incarnation:
-                    def save_goal_state(incarnation, xml_text):
-                        file_name = GOAL_STATE_FILE_NAME.format(incarnation)
-                        goal_state_file = os.path.join(conf.get_lib_dir(), file_name)
-                        self.save_cache(goal_state_file, xml_text)
-
-                    self.goal_state_flusher.flush(datetime.utcnow())
-                    self._set_goal_state(new_goal_state)
-                    save_goal_state(new_goal_state.incarnation, new_goal_state_xml)
-                    self.update_hosting_env(new_goal_state)
-                    self.update_shared_conf(new_goal_state)
-                    self.update_certs(new_goal_state)
-                    self.update_ext_conf(new_goal_state)
-                    self.update_remote_access_conf(new_goal_state)
-                    self.save_cache(os.path.join(conf.get_lib_dir(), INCARNATION_FILE_NAME), new_goal_state.incarnation)
-
-                    update_host_plugin()
+                if self._goal_state is None or refresh_type == WireClient._UpdateType.GoalStateForced:
+                    new_goal_state = GoalState.fetch_full_goal_state(self)
+                else:
+                    new_goal_state = GoalState.fetch_full_goal_state_if_incarnation_different_than(self, self._goal_state.incarnation)
+
+                if new_goal_state is not None:
+                    self._goal_state = new_goal_state
+                    self._save_goal_state()
+                    self._update_host_plugin(new_goal_state.container_id, new_goal_state.role_config_name)
+
                 return
 
             except IOError as e:
@@ -826,35 +745,37 @@ class WireClient(object):
 
         raise ProtocolError("Exceeded max retry updating goal state")
 
-    def _set_goal_state(self, new_goal_state):
-        if new_goal_state is None:
-            logger.warn("Setting empty Goal State object!")
-        self._goal_state = new_goal_state
+    def _update_host_plugin(self, container_id, role_config_name):
+        if self._host_plugin is not None:
+            self._host_plugin.update_container_id(container_id)
+            self._host_plugin.update_role_config_name(role_config_name)
 
-    def _set_hosting_env(self, new_hosting_env):
-        if new_hosting_env is None:
-            logger.warn("Setting empty Hosting Environment object!")
-        self._hosting_env = new_hosting_env
+    def _save_goal_state(self):
+        try:
+            self.goal_state_flusher.flush(datetime.utcnow())
 
-    def _set_shared_conf(self, new_shared_config):
-        if new_shared_config is None:
-            logger.warn("Setting empty Shared Config object!")
-        self._shared_conf = new_shared_config
+        except Exception as e:
+            logger.warn("Failed to save the previous goal state to the history folder: {0}", ustr(e))
+
+        try:
+            local_file = os.path.join(conf.get_lib_dir(), INCARNATION_FILE_NAME)
+            self.save_cache(local_file, self._goal_state.incarnation)
 
-    def _set_certs(self, new_certs):
-        if new_certs is None:
-            logger.warn("Setting empty Certs object!")
-        self._certs = new_certs
+            def save_if_not_none(goal_state_property, file_name):
+                file_path = os.path.join(conf.get_lib_dir(), file_name)
 
-    def _set_ext_conf(self, new_ext_conf):
-        if new_ext_conf is None:
-            logger.warn("Setting empty Extension Config object!")
-        self._ext_conf = new_ext_conf
+                if goal_state_property is not None and goal_state_property.xml_text is not None:
+                    self.save_cache(file_path, goal_state_property.xml_text)
 
-    def _set_remote_access(self, new_remote_access):
-        if new_remote_access is None:
-            logger.warn("Setting empty Remote Access object!")
-        self._remote_access = new_remote_access
+            # NOTE: Certificates are saved in Certificate.__init__
+            save_if_not_none(self._goal_state, GOAL_STATE_FILE_NAME.format(self._goal_state.incarnation))
+            save_if_not_none(self._goal_state.hosting_env, HOSTING_ENV_FILE_NAME)
+            save_if_not_none(self._goal_state.shared_conf, SHARED_CONF_FILE_NAME)
+            save_if_not_none(self._goal_state.ext_conf, EXT_CONF_FILE_NAME.format(self._goal_state.incarnation))
+            save_if_not_none(self._goal_state.remote_access, REMOTE_ACCESS_FILE_NAME.format(self._goal_state.incarnation))
+
+        except Exception as e:
+            logger.warn("Failed to save the goal state to disk: {0}", ustr(e))
 
     def _set_host_plugin(self, new_host_plugin):
         if new_host_plugin is None:
@@ -867,29 +788,30 @@ class WireClient(object):
         return self._goal_state
 
     def get_hosting_env(self):
-        if self._hosting_env is None:
+        if self._goal_state is None:
             raise ProtocolError("Trying to fetch Hosting Environment before initialization!")
-        return self._hosting_env
+        return self._goal_state.hosting_env
 
     def get_shared_conf(self):
-        if self._shared_conf is None:
+        if self._goal_state is None:
             raise ProtocolError("Trying to fetch Shared Conf before initialization!")
-        return self._shared_conf
+        return self._goal_state.shared_conf
 
     def get_certs(self):
-        # This property can be None in the GoalState, so not returning a ProtocolError here
-        return self._certs
+        if self._goal_state is None:
+            raise ProtocolError("Trying to fetch Certificates before initialization!")
+        return self._goal_state.certs
 
     def get_ext_conf(self):
-        if self._ext_conf is None:
-            if self.get_goal_state().ext_uri is None:
-                self._set_ext_conf(ExtensionsConfig(None))
-            else:
-                raise ProtocolError("Trying to fetch Extension Conf before initialization!")
-        return self._ext_conf
+        if self._goal_state is None:
+            raise ProtocolError("Trying to fetch Extension Conf before initialization!")
+        return self._goal_state.ext_conf
 
-    def get_ext_manifest(self, ext_handler, goal_state):
-        local_file = MANIFEST_FILE_NAME.format(ext_handler.name, goal_state.incarnation)
+    def get_ext_manifest(self, ext_handler):
+        if self._goal_state is None:
+            raise ProtocolError("Trying to fetch Extension Manifest before initialization!")
+
+        local_file = MANIFEST_FILE_NAME.format(ext_handler.name, self.get_goal_state().incarnation)
         local_file = os.path.join(conf.get_lib_dir(), local_file)
 
         try:
@@ -899,6 +821,11 @@ class WireClient(object):
         except Exception as e:
             raise ExtensionDownloadError("Failed to retrieve extension manifest. Error: {0}".format(ustr(e)))
 
+    def get_remote_access(self):
+        if self._goal_state is None:
+            raise ProtocolError("Trying to fetch Remote Access before initialization!")
+        return self._goal_state.remote_access
+
     def get_gafamily_manifest(self, vmagent_manifest, goal_state):
         local_file = MANIFEST_FILE_NAME.format(vmagent_manifest.family, goal_state.incarnation)
         local_file = os.path.join(conf.get_lib_dir(), local_file)
@@ -1285,365 +1212,6 @@ class VersionInfo(object):
         return self.supported
 
 
-class GoalState(object):
-    def __init__(self, xml_text):
-        if xml_text is None:
-            raise ValueError("GoalState.xml is None")
-        logger.verbose("Load GoalState.xml")
-        self.incarnation = None
-        self.expected_state = None
-        self.hosting_env_uri = None
-        self.shared_conf_uri = None
-        self.remote_access_uri = None
-        self.certs_uri = None
-        self.ext_uri = None
-        self.role_instance_id = None
-        self.role_config_name = None
-        self.container_id = None
-        self.load_balancer_probe_port = None
-        self.xml_text = None
-        self._parse(xml_text)
-
-    def _parse(self, xml_text):
-        self.xml_text = xml_text
-        xml_doc = parse_doc(xml_text)
-        self.incarnation = findtext(xml_doc, "Incarnation")
-        self.expected_state = findtext(xml_doc, "ExpectedState")
-        self.hosting_env_uri = findtext(xml_doc, "HostingEnvironmentConfig")
-        self.shared_conf_uri = findtext(xml_doc, "SharedConfig")
-        self.certs_uri = findtext(xml_doc, "Certificates")
-        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_config_name = findtext(role_config, "ConfigName")
-        container = find(xml_doc, "Container")
-        self.container_id = findtext(container, "ContainerId")
-        os.environ[CONTAINER_ID_ENV_VARIABLE] = self.container_id
-        self.remote_access_uri = findtext(container, "RemoteAccessInfo")
-        lbprobe_ports = find(xml_doc, "LBProbePorts")
-        self.load_balancer_probe_port = findtext(lbprobe_ports, "Port")
-
-
-class HostingEnv(object):
-    def __init__(self, xml_text):
-        if xml_text is None:
-            raise ValueError("HostingEnvironmentConfig.xml is None")
-        logger.verbose("Load HostingEnvironmentConfig.xml")
-        self.vm_name = None
-        self.role_name = None
-        self.deployment_name = None
-        self.xml_text = None
-        self._parse(xml_text)
-
-    def _parse(self, xml_text):
-        self.xml_text = xml_text
-        xml_doc = parse_doc(xml_text)
-        incarnation = find(xml_doc, "Incarnation")
-        self.vm_name = getattrib(incarnation, "instance")
-        role = find(xml_doc, "Role")
-        self.role_name = getattrib(role, "name")
-        deployment = find(xml_doc, "Deployment")
-        self.deployment_name = getattrib(deployment, "name")
-
-
-class SharedConfig(object):
-    def __init__(self, xml_text):
-        logger.verbose("Load SharedConfig.xml")
-        self.xml_text = xml_text
-
-
-class RemoteAccess(object):
-    """
-    Object containing information about user accounts
-    """
-
-    #
-    # <RemoteAccess>
-    #   <Version/>
-    #   <Incarnation/>
-    #    <Users>
-    #       <User>
-    #         <Name/>
-    #         <Password/>
-    #         <Expiration/>
-    #       </User>
-    #     </Users>
-    #   </RemoteAccess>
-    #
-
-    def __init__(self, xml_text):
-        logger.verbose("Load RemoteAccess.xml")
-        self.version = None
-        self.incarnation = None
-        self.user_list = RemoteAccessUsersList()
-
-        self.xml_text = None
-        self._parse(xml_text)
-
-    def _parse(self, xml_text):
-        """
-        Parse xml document containing user account information
-        """
-        if xml_text is None or len(xml_text) == 0:
-            return None
-        self.xml_text = xml_text
-        xml_doc = parse_doc(xml_text)
-        self.incarnation = findtext(xml_doc, "Incarnation")
-        self.version = findtext(xml_doc, "Version")
-        user_collection = find(xml_doc, "Users")
-        users = findall(user_collection, "User")
-
-        for user in users:
-            remote_access_user = self._parse_user(user)
-            self.user_list.users.append(remote_access_user)
-
-    def _parse_user(self, user):
-        name = findtext(user, "Name")
-        encrypted_password = findtext(user, "Password")
-        expiration = findtext(user, "Expiration")
-        remote_access_user = RemoteAccessUser(name, encrypted_password, expiration)
-        return remote_access_user
-
-
-class UserAccount(object):
-    """
-    Stores information about single user account
-    """
-
-    def __init__(self):
-        self.Name = None
-        self.EncryptedPassword = None
-        self.Password = None
-        self.Expiration = None
-        self.Groups = []
-
-
-class Certificates(object):
-    """
-    Object containing certificates of host and provisioned user.
-    """
-
-    def __init__(self, client, xml_text):
-        logger.verbose("Load Certificates.xml")
-        self.client = client
-        self.cert_list = CertList()
-        self._parse(xml_text)
-
-    def _parse(self, xml_text):
-        """
-        Parse multiple certificates into seperate files.
-        """
-        xml_doc = parse_doc(xml_text)
-        data = findtext(xml_doc, "Data")
-        if data is None:
-            return
-
-        # if the certificates format is not Pkcs7BlobWithPfxContents do not parse it
-        certificateFormat = findtext(xml_doc, "Format")
-        if certificateFormat and certificateFormat != "Pkcs7BlobWithPfxContents":
-            logger.warn("The Format is not Pkcs7BlobWithPfxContents. Format is " + certificateFormat)
-            return
-
-        cryptutil = CryptUtil(conf.get_openssl_cmd())
-        p7m_file = os.path.join(conf.get_lib_dir(), P7M_FILE_NAME)
-        p7m = ("MIME-Version:1.0\n"
-               "Content-Disposition: attachment; filename=\"{0}\"\n"
-               "Content-Type: application/x-pkcs7-mime; name=\"{1}\"\n"
-               "Content-Transfer-Encoding: base64\n"
-               "\n"
-               "{2}").format(p7m_file, p7m_file, data)
-
-        self.client.save_cache(p7m_file, p7m)
-
-        trans_prv_file = os.path.join(conf.get_lib_dir(),
-                                      TRANSPORT_PRV_FILE_NAME)
-        trans_cert_file = os.path.join(conf.get_lib_dir(),
-                                       TRANSPORT_CERT_FILE_NAME)
-        pem_file = os.path.join(conf.get_lib_dir(), PEM_FILE_NAME)
-        # decrypt certificates
-        cryptutil.decrypt_p7m(p7m_file, trans_prv_file, trans_cert_file,
-                              pem_file)
-
-        # The parsing process use public key to match prv and crt.
-        buf = []
-        begin_crt = False
-        begin_prv = False
-        prvs = {}
-        thumbprints = {}
-        index = 0
-        v1_cert_list = []
-        with open(pem_file) as pem:
-            for line in pem.readlines():
-                buf.append(line)
-                if re.match(r'[-]+BEGIN.*KEY[-]+', line):
-                    begin_prv = True
-                elif re.match(r'[-]+BEGIN.*CERTIFICATE[-]+', line):
-                    begin_crt = True
-                elif re.match(r'[-]+END.*KEY[-]+', line):
-                    tmp_file = self.write_to_tmp_file(index, 'prv', buf)
-                    pub = cryptutil.get_pubkey_from_prv(tmp_file)
-                    prvs[pub] = tmp_file
-                    buf = []
-                    index += 1
-                    begin_prv = False
-                elif re.match(r'[-]+END.*CERTIFICATE[-]+', line):
-                    tmp_file = self.write_to_tmp_file(index, 'crt', buf)
-                    pub = cryptutil.get_pubkey_from_crt(tmp_file)
-                    thumbprint = cryptutil.get_thumbprint_from_crt(tmp_file)
-                    thumbprints[pub] = thumbprint
-                    # Rename crt with thumbprint as the file name
-                    crt = "{0}.crt".format(thumbprint)
-                    v1_cert_list.append({
-                        "name": None,
-                        "thumbprint": thumbprint
-                    })
-                    os.rename(tmp_file, os.path.join(conf.get_lib_dir(), crt))
-                    buf = []
-                    index += 1
-                    begin_crt = False
-
-        # Rename prv key with thumbprint as the file name
-        for pubkey in prvs:
-            thumbprint = thumbprints[pubkey]
-            if thumbprint:
-                tmp_file = prvs[pubkey]
-                prv = "{0}.prv".format(thumbprint)
-                os.rename(tmp_file, os.path.join(conf.get_lib_dir(), prv))
-                logger.info("Found private key matching thumbprint {0}".format(thumbprint))
-            else:
-                # Since private key has *no* matching certificate,
-                # it will not be named correctly
-                logger.warn("Found NO matching cert/thumbprint for private key!")
-
-        # Log if any certificates were found without matching private keys
-        # This can happen (rarely), and is useful to know for debugging
-        for pubkey in thumbprints:
-            if not pubkey in prvs:
-                msg = "Certificate with thumbprint {0} has no matching private key."
-                logger.info(msg.format(thumbprints[pubkey]))
-
-        for v1_cert in v1_cert_list:
-            cert = Cert()
-            set_properties("certs", cert, v1_cert)
-            self.cert_list.certificates.append(cert)
-
-    def write_to_tmp_file(self, index, suffix, buf):
-        file_name = os.path.join(conf.get_lib_dir(),
-                                 "{0}.{1}".format(index, suffix))
-        self.client.save_cache(file_name, "".join(buf))
-        return file_name
-
-
-class ExtensionsConfig(object):
-    def __init__(self, xml_text):
-        logger.verbose("Load ExtensionsConfig.xml")
-        self.ext_handlers = ExtHandlerList()
-        self.vmagent_manifests = VMAgentManifestList()
-        self.status_upload_blob = None
-        self.status_upload_blob_type = None
-        self.artifacts_profile_blob = None
-        if xml_text is not None:
-            self._parse(xml_text)
-
-    def _parse(self, xml_text):
-        xml_doc = parse_doc(xml_text)
-
-        ga_families_list = find(xml_doc, "GAFamilies")
-        ga_families = findall(ga_families_list, "GAFamily")
-
-        for ga_family in ga_families:
-            family = findtext(ga_family, "Name")
-            uris_list = find(ga_family, "Uris")
-            uris = findall(uris_list, "Uri")
-            manifest = VMAgentManifest()
-            manifest.family = family
-            for uri in uris:
-                manifestUri = VMAgentManifestUri(uri=gettext(uri))
-                manifest.versionsManifestUris.append(manifestUri)
-            self.vmagent_manifests.vmAgentManifests.append(manifest)
-
-        plugins_list = find(xml_doc, "Plugins")
-        plugins = findall(plugins_list, "Plugin")
-        plugin_settings_list = find(xml_doc, "PluginSettings")
-        plugin_settings = findall(plugin_settings_list, "Plugin")
-
-        for plugin in plugins:
-            ext_handler = self._parse_plugin(plugin)
-            self.ext_handlers.extHandlers.append(ext_handler)
-            self._parse_plugin_settings(ext_handler, plugin_settings)
-
-        self.status_upload_blob = findtext(xml_doc, "StatusUploadBlob")
-        self.artifacts_profile_blob = findtext(xml_doc, "InVMArtifactsProfileBlob")
-
-        status_upload_node = find(xml_doc, "StatusUploadBlob")
-        self.status_upload_blob_type = getattrib(status_upload_node,
-                                                 "statusBlobType")
-        logger.verbose("Extension config shows status blob type as [{0}]",
-                       self.status_upload_blob_type)
-
-    def _parse_plugin(self, plugin):
-        ext_handler = ExtHandler()
-        ext_handler.name = getattrib(plugin, "name")
-        ext_handler.properties.version = getattrib(plugin, "version")
-        ext_handler.properties.state = getattrib(plugin, "state")
-
-        location = getattrib(plugin, "location")
-        failover_location = getattrib(plugin, "failoverlocation")
-        for uri in [location, failover_location]:
-            version_uri = ExtHandlerVersionUri()
-            version_uri.uri = uri
-            ext_handler.versionUris.append(version_uri)
-        return ext_handler
-
-    def _parse_plugin_settings(self, ext_handler, plugin_settings):
-        if plugin_settings is None:
-            return
-
-        name = ext_handler.name
-        version = ext_handler.properties.version
-        settings = [x for x in plugin_settings \
-                    if getattrib(x, "name") == name and \
-                    getattrib(x, "version") == version]
-
-        if settings is None or len(settings) == 0:
-            return
-
-        runtime_settings = None
-        runtime_settings_node = find(settings[0], "RuntimeSettings")
-        seqNo = getattrib(runtime_settings_node, "seqNo")
-        runtime_settings_str = gettext(runtime_settings_node)
-        try:
-            runtime_settings = json.loads(runtime_settings_str)
-        except ValueError as e:
-            logger.error("Invalid extension settings")
-            return
-
-        depends_on_level = 0
-        depends_on_node = find(settings[0], "DependsOn")
-        if depends_on_node != None:
-            try:
-                depends_on_level = int(getattrib(depends_on_node, "dependencyLevel"))
-            except (ValueError, TypeError):
-                logger.warn("Could not parse dependencyLevel for handler {0}. Setting it to 0".format(name))
-                depends_on_level = 0
-
-        for plugin_settings_list in runtime_settings["runtimeSettings"]:
-            handler_settings = plugin_settings_list["handlerSettings"]
-            ext = Extension()
-            # There is no "extension name" in wire protocol.
-            # Put
-            ext.name = ext_handler.name
-            ext.sequenceNumber = seqNo
-            ext.publicSettings = handler_settings.get("publicSettings")
-            ext.protectedSettings = handler_settings.get("protectedSettings")
-            ext.dependencyLevel = depends_on_level
-            thumbprint = handler_settings.get(
-                "protectedSettingsCertThumbprint")
-            ext.certificateThumbprint = thumbprint
-            ext_handler.properties.extensions.append(ext)
-
-
 class ExtensionManifest(object):
     def __init__(self, xml_text):
         if xml_text is None:
diff --git a/azurelinuxagent/daemon/main.py b/azurelinuxagent/daemon/main.py
index d328daf9..5ba1a476 100644
--- a/azurelinuxagent/daemon/main.py
+++ b/azurelinuxagent/daemon/main.py
@@ -30,7 +30,7 @@ from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator
 from azurelinuxagent.common.event import add_event, WALAEventOperation
 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.util import get_protocol_util
 from azurelinuxagent.common.protocol.wire import WireProtocol
 from azurelinuxagent.common.rdma import setup_rdma_device
 from azurelinuxagent.common.version import AGENT_NAME, AGENT_LONG_NAME, \
diff --git a/azurelinuxagent/ga/env.py b/azurelinuxagent/ga/env.py
index 371fa38b..9ea16155 100644
--- a/azurelinuxagent/ga/env.py
+++ b/azurelinuxagent/ga/env.py
@@ -31,7 +31,7 @@ from azurelinuxagent.common.dhcp import get_dhcp_handler
 from azurelinuxagent.common.event import add_periodic, WALAEventOperation
 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.util import get_protocol_util
 from azurelinuxagent.common.utils.archive import StateArchiver
 from azurelinuxagent.common.version import AGENT_NAME, CURRENT_VERSION
 
diff --git a/azurelinuxagent/ga/monitor.py b/azurelinuxagent/ga/monitor.py
index ae208925..e0bfe8d4 100644
--- a/azurelinuxagent/ga/monitor.py
+++ b/azurelinuxagent/ga/monitor.py
@@ -34,7 +34,7 @@ from azurelinuxagent.common.event import add_event, EventLogger, get_container_i
 from azurelinuxagent.common.exception import EventError, HttpError, OSUtilError, 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.util import get_protocol_util
 from azurelinuxagent.common.protocol.healthservice import HealthService
 from azurelinuxagent.common.protocol.imds import get_imds_client
 from azurelinuxagent.common.telemetryevent import TelemetryEvent, TelemetryEventList, TelemetryEventParam
diff --git a/azurelinuxagent/ga/update.py b/azurelinuxagent/ga/update.py
index ccdc40b4..f44b294e 100644
--- a/azurelinuxagent/ga/update.py
+++ b/azurelinuxagent/ga/update.py
@@ -50,7 +50,7 @@ from azurelinuxagent.common.exception import ProtocolError, \
                                             UpdateError
 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.util import get_protocol_util
 from azurelinuxagent.common.protocol.hostplugin import HostPluginProtocol
 from azurelinuxagent.common.protocol.wire import WireProtocol
 from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
diff --git a/azurelinuxagent/pa/deprovision/default.py b/azurelinuxagent/pa/deprovision/default.py
index 5205a966..eb783b3e 100644
--- a/azurelinuxagent/pa/deprovision/default.py
+++ b/azurelinuxagent/pa/deprovision/default.py
@@ -30,7 +30,7 @@ from azurelinuxagent.common import version
 
 from azurelinuxagent.common.exception import ProtocolError
 from azurelinuxagent.common.osutil import get_osutil
-from azurelinuxagent.common.protocol import get_protocol_util
+from azurelinuxagent.common.protocol.util import get_protocol_util
 from azurelinuxagent.ga.exthandlers import HANDLER_NAME_PATTERN
 
 
diff --git a/azurelinuxagent/pa/provision/cloudinit.py b/azurelinuxagent/pa/provision/cloudinit.py
index e4455685..aff2be06 100644
--- a/azurelinuxagent/pa/provision/cloudinit.py
+++ b/azurelinuxagent/pa/provision/cloudinit.py
@@ -32,7 +32,7 @@ import azurelinuxagent.common.utils.shellutil as shellutil
 from azurelinuxagent.common.event import elapsed_milliseconds, WALAEventOperation
 from azurelinuxagent.common.exception import ProvisionError, ProtocolError
 from azurelinuxagent.common.future import ustr
-from azurelinuxagent.common.protocol import OVF_FILE_NAME
+from azurelinuxagent.common.protocol.util import OVF_FILE_NAME
 from azurelinuxagent.common.protocol.ovfenv import OvfEnv
 from azurelinuxagent.pa.provision.default import ProvisionHandler
 
diff --git a/azurelinuxagent/pa/provision/default.py b/azurelinuxagent/pa/provision/default.py
index 6efdaf9e..7b066dcc 100644
--- a/azurelinuxagent/pa/provision/default.py
+++ b/azurelinuxagent/pa/provision/default.py
@@ -38,7 +38,7 @@ from azurelinuxagent.common.exception import ProvisionError, ProtocolError, \
     OSUtilError
 from azurelinuxagent.common.osutil import get_osutil
 from azurelinuxagent.common.protocol.restapi import ProvisionStatus
-from azurelinuxagent.common.protocol import get_protocol_util
+from azurelinuxagent.common.protocol.util import get_protocol_util
 from azurelinuxagent.common.version import AGENT_NAME
 
 CUSTOM_DATA_FILE = "CustomData"


Test patch:
diff --git a/tests/common/test_event.py b/tests/common/test_event.py
index 8f71f1c1..1d4f00cc 100644
--- a/tests/common/test_event.py
+++ b/tests/common/test_event.py
@@ -32,7 +32,8 @@ from azurelinuxagent.common.utils import fileutil
 from azurelinuxagent.common.utils.extensionprocessutil import read_output
 from azurelinuxagent.common.version import CURRENT_AGENT, CURRENT_VERSION
 from azurelinuxagent.ga.monitor import MonitorHandler
-from tests.tools import AgentTestCase, data_dir, load_data, Mock, patch
+from tests.protocol import mockwiredata, mock_wire_protocol
+from tests.tools import AgentTestCase, data_dir, Mock, patch
 
 
 class TestEvent(AgentTestCase):
@@ -59,29 +60,28 @@ class TestEvent(AgentTestCase):
                             '{{"value": "{0}", "name": "ContainerId"}}'.format(
                                 os.environ[event.CONTAINER_ID_ENV_VARIABLE]) in data)
 
-            # Container id is set as an environment variable when parsing the goal state
-            xml_text = load_data("wire/goal_state.xml")
-            goal_state = GoalState(xml_text)
+            with mock_wire_protocol.create(mockwiredata.DATA_FILE) as protocol:
+                goal_state = protocol.client._goal_state
 
-            container_id = goal_state.container_id
-            event.add_event(name='dummy_name')
-            data = fileutil.read_file(tmp_file)
-            self.assertTrue('{{"name": "ContainerId", "value": "{0}"}}'.format(container_id) in data or
-                            '{{"value": "{0}", "name": "ContainerId"}}'.format(container_id), data)
+                # Container id is set as an environment variable when parsing the goal state
+                container_id = goal_state.container_id
+                event.add_event(name='dummy_name')
+                data = fileutil.read_file(tmp_file)
+                self.assertTrue('{{"name": "ContainerId", "value": "{0}"}}'.format(container_id) in data or
+                                '{{"value": "{0}", "name": "ContainerId"}}'.format(container_id), data)
 
-            # Container id is updated as the goal state changes, both in telemetry event and in environment variables
-            new_container_id = "z6d5526c-5ac2-4200-b6e2-56f2b70c5ab2"
-            xml_text = load_data("wire/goal_state.xml")
-            xml_text_updated = xml_text.replace("c6d5526c-5ac2-4200-b6e2-56f2b70c5ab2", new_container_id)
-            goal_state = GoalState(xml_text_updated)
+                # Container id is updated as the goal state changes, both in telemetry event and in environment variables
+                new_container_id = "z6d5526c-5ac2-4200-b6e2-56f2b70c5ab2"
+                protocol.mock_wire_data.set_container_id(new_container_id)
+                protocol.client.update_goal_state()
 
-            event.add_event(name='dummy_name')
-            data = fileutil.read_file(tmp_file)
+                event.add_event(name='dummy_name')
+                data = fileutil.read_file(tmp_file)
 
-            # Assert both the environment variable and telemetry event got updated
-            self.assertEquals(os.environ[event.CONTAINER_ID_ENV_VARIABLE], new_container_id)
-            self.assertTrue('{{"name": "ContainerId", "value": "{0}"}}'.format(new_container_id) in data or
-                            '{{"value": "{0}", "name": "ContainerId"}}'.format(new_container_id), data)
+                # Assert both the environment variable and telemetry event got updated
+                self.assertEquals(os.environ[event.CONTAINER_ID_ENV_VARIABLE], new_container_id)
+                self.assertTrue('{{"name": "ContainerId", "value": "{0}"}}'.format(new_container_id) in data or
+                                '{{"value": "{0}", "name": "ContainerId"}}'.format(new_container_id), data)
 
         os.environ.pop(event.CONTAINER_ID_ENV_VARIABLE)
 
diff --git a/tests/data/wire/ext_conf.xml b/tests/data/wire/ext_conf.xml
index 3bafcef1..8f205fb8 100644
--- a/tests/data/wire/ext_conf.xml
+++ b/tests/data/wire/ext_conf.xml
@@ -3,19 +3,19 @@
     <GAFamily>
       <Name>Prod</Name>
       <Uris>
-          <Uri>http://manifest_of_ga.xml</Uri>
+          <Uri>http://mock-goal-state/manifest_of_ga.xml</Uri>
       </Uris>
     </GAFamily>
     <GAFamily>
       <Name>Test</Name>
       <Uris>
-          <Uri>http://manifest_of_ga.xml</Uri>
+          <Uri>http://mock-goal-state/manifest_of_ga.xml</Uri>
         </Uris>
     </GAFamily>
   </GAFamilies>
 </GuestAgentExtension>
 <Plugins>
-  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0" location="http://rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="false" failoverlocation="http://rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
+  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0" location="http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="false" failoverlocation="http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
 </Plugins>
 <PluginSettings>
   <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0">
diff --git a/tests/data/wire/ext_conf_autoupgrade.xml b/tests/data/wire/ext_conf_autoupgrade.xml
index 67e95c3d..313a1af3 100644
--- a/tests/data/wire/ext_conf_autoupgrade.xml
+++ b/tests/data/wire/ext_conf_autoupgrade.xml
@@ -3,21 +3,21 @@
     <GAFamily>
       <Name>Win8</Name>
       <Uris>
-        <Uri>http://rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
       </Uris>
     </GAFamily>
     <GAFamily>
       <Name>Win7</Name>
       <Uris>
-        <Uri>http://rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
       </Uris>
     </GAFamily>
   </GAFamilies>
 </GuestAgentExtension>
 <Plugins>
-  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0" location="http://rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="true" failoverlocation="http://rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
+  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0" location="http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="true" failoverlocation="http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
 </Plugins>
 <PluginSettings>
   <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0">
diff --git a/tests/data/wire/ext_conf_autoupgrade_internalversion.xml b/tests/data/wire/ext_conf_autoupgrade_internalversion.xml
index d918482f..97f840c1 100644
--- a/tests/data/wire/ext_conf_autoupgrade_internalversion.xml
+++ b/tests/data/wire/ext_conf_autoupgrade_internalversion.xml
@@ -3,21 +3,21 @@
     <GAFamily>
       <Name>Win8</Name>
       <Uris>
-        <Uri>http://rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
       </Uris>
     </GAFamily>
     <GAFamily>
       <Name>Win7</Name>
       <Uris>
-        <Uri>http://rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
       </Uris>
     </GAFamily>
   </GAFamilies>
 </GuestAgentExtension>
 <Plugins>
-  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.3.0" location="http://rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="true" failoverlocation="http://rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
+  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.3.0" location="http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="true" failoverlocation="http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
 </Plugins>
 <PluginSettings>
   <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.3.0">
diff --git a/tests/data/wire/ext_conf_internalversion.xml b/tests/data/wire/ext_conf_internalversion.xml
index d918482f..97f840c1 100644
--- a/tests/data/wire/ext_conf_internalversion.xml
+++ b/tests/data/wire/ext_conf_internalversion.xml
@@ -3,21 +3,21 @@
     <GAFamily>
       <Name>Win8</Name>
       <Uris>
-        <Uri>http://rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
       </Uris>
     </GAFamily>
     <GAFamily>
       <Name>Win7</Name>
       <Uris>
-        <Uri>http://rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
       </Uris>
     </GAFamily>
   </GAFamilies>
 </GuestAgentExtension>
 <Plugins>
-  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.3.0" location="http://rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="true" failoverlocation="http://rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
+  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.3.0" location="http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="true" failoverlocation="http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
 </Plugins>
 <PluginSettings>
   <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.3.0">
diff --git a/tests/data/wire/ext_conf_missing_family.xml b/tests/data/wire/ext_conf_missing_family.xml
index 40100779..8f00bc9c 100644
--- a/tests/data/wire/ext_conf_missing_family.xml
+++ b/tests/data/wire/ext_conf_missing_family.xml
@@ -10,20 +10,20 @@
             <GAFamily>
                 <Name>Test</Name>
                 <Uris>
-                    <Uri>https://rdfepirv2bl2prdstr01.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://rdfepirv2bl2prdstr02.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://rdfepirv2bl2prdstr03.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://rdfepirv2bl2prdstr04.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://rdfepirv2bl3prdstr01.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://rdfepirv2bl3prdstr02.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://rdfepirv2bl3prdstr03.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://zrdfepirv2bl4prdstr01.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://zrdfepirv2bl4prdstr03.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://zrdfepirv2bl5prdstr02.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://zrdfepirv2bl5prdstr04.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://zrdfepirv2bl5prdstr06.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://zrdfepirv2bl5prdstr09a.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
-                    <Uri>https://zrdfepirv2bl6prdstr02a.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/rdfepirv2bl2prdstr01.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/rdfepirv2bl2prdstr02.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/rdfepirv2bl2prdstr03.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/rdfepirv2bl2prdstr04.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/rdfepirv2bl3prdstr01.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/rdfepirv2bl3prdstr02.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/rdfepirv2bl3prdstr03.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/zrdfepirv2bl4prdstr01.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/zrdfepirv2bl4prdstr03.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/zrdfepirv2bl5prdstr02.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/zrdfepirv2bl5prdstr04.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/zrdfepirv2bl5prdstr06.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/zrdfepirv2bl5prdstr09a.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
+                    <Uri>https://mock-goal-state/zrdfepirv2bl6prdstr02a.blob.core.windows.net/7d89d439b79f4452950452399add2c90/Microsoft.OSTCLinuxAgent_Test_useast_manifest.xml</Uri>
                 </Uris>
             </GAFamily>
         </GAFamilies>
diff --git a/tests/data/wire/ext_conf_multiple_extensions.xml b/tests/data/wire/ext_conf_multiple_extensions.xml
index c68f35eb..b6e082a2 100644
--- a/tests/data/wire/ext_conf_multiple_extensions.xml
+++ b/tests/data/wire/ext_conf_multiple_extensions.xml
@@ -4,42 +4,42 @@
             <GAFamily>
                 <Name>Prod</Name>
                 <Uris>
-                    <Uri>http://manifest_of_ga.xml</Uri>
+                    <Uri>http://mock-goal-state/manifest_of_ga.xml</Uri>
                 </Uris>
             </GAFamily>
             <GAFamily>
                 <Name>Test</Name>
                 <Uris>
-                    <Uri>http://manifest_of_ga.xml</Uri>
+                    <Uri>http://mock-goal-state/manifest_of_ga.xml</Uri>
                 </Uris>
             </GAFamily>
         </GAFamilies>
     </GuestAgentExtension>
     <Plugins>
         <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0"
-                location="http://rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml"
+                location="http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml"
                 config="" state="enabled" autoUpgrade="false"
-                failoverlocation="http://rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml"
+                failoverlocation="http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml"
                 runAsStartupTask="false" isJson="true"/>
         <Plugin name="Microsoft.Powershell.ExampleExtension" version="1.0.0"
-                location="https://zrdfepirv2sn3prdstr06.blob.core.windows.net/3055d769a0ed485d8c5ea8a7cd24073e/Microsoft.Powershell_ExampleExtension_ussouth_manifest.xml"
+                location="https://mock-goal-state/zrdfepirv2sn3prdstr06.blob.core.windows.net/3055d769a0ed485d8c5ea8a7cd24073e/Microsoft.Powershell_ExampleExtension_ussouth_manifest.xml"
                 state="enabled" autoUpgrade="true"
-                failoverlocation="https://zrdfepirv2sn4prdstr02a.blob.core.windows.net/3055d769a0ed485d8c5ea8a7cd24073e/Microsoft.Powershell_ExampleExtension_ussouth_manifest.xml"
+                failoverlocation="https://mock-goal-state/zrdfepirv2sn4prdstr02a.blob.core.windows.net/3055d769a0ed485d8c5ea8a7cd24073e/Microsoft.Powershell_ExampleExtension_ussouth_manifest.xml"
                 runAsStartupTask="false" isJson="true" useExactVersion="false"/>
         <Plugin name="Microsoft.EnterpriseCloud.Monitoring.ExampleHandlerLinux" version="1.0.0"
-                location="https://rdfepirv2sn3prdstr01.blob.core.windows.net/8b23413e0ce64aa89889a295b6122112/Microsoft.EnterpriseCloud.Monitoring_ExampleHandlerLinux_ussouth_manifest.xml"
+                location="https://mock-goal-state/rdfepirv2sn3prdstr01.blob.core.windows.net/8b23413e0ce64aa89889a295b6122112/Microsoft.EnterpriseCloud.Monitoring_ExampleHandlerLinux_ussouth_manifest.xml"
                 state="enabled" autoUpgrade="true"
-                failoverlocation="https://ardfepirv2sn3prdstr05.blob.core.windows.net/8b23413e0ce64aa89889a295b6122112/Microsoft.EnterpriseCloud.Monitoring_ExampleHandlerLinux_ussouth_manifest.xml"
+                failoverlocation="https://mock-goal-state/ardfepirv2sn3prdstr05.blob.core.windows.net/8b23413e0ce64aa89889a295b6122112/Microsoft.EnterpriseCloud.Monitoring_ExampleHandlerLinux_ussouth_manifest.xml"
                 runAsStartupTask="false" isJson="true" useExactVersion="false"/>
         <Plugin name="Microsoft.CPlat.Core.ExampleExtensionLinux" version="1.0.0"
-                location="https://zrdfepirv2sn4prdstr09a.blob.core.windows.net/f72653efd9e349ed9842c8b99e4c1712/Microsoft.CPlat.Core_ExampleExtensionLinux_ussouth_manifest.xml"
+                location="https://mock-goal-state/zrdfepirv2sn4prdstr09a.blob.core.windows.net/f72653efd9e349ed9842c8b99e4c1712/Microsoft.CPlat.Core_ExampleExtensionLinux_ussouth_manifest.xml"
                 state="enabled" autoUpgrade="true"
-                failoverlocation="https://ardfepirv2sn3prdstr05.blob.core.windows.net/f72653efd9e349ed9842c8b99e4c1712/Microsoft.CPlat.Core_ExampleExtensionLinux_ussouth_manifest.xml"
+                failoverlocation="https://mock-goal-state/ardfepirv2sn3prdstr05.blob.core.windows.net/f72653efd9e349ed9842c8b99e4c1712/Microsoft.CPlat.Core_ExampleExtensionLinux_ussouth_manifest.xml"
                 runAsStartupTask="false" isJson="true" useExactVersion="false"/>
         <Plugin name="Microsoft.OSTCExtensions.Edp.ExampleExtensionLinuxInTest" version="1.0.0"
-                location="https://zrdfepirv2sn6prdstr05a.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions.Edp_ExampleExtensionLinuxInTest_ussouth_manifest.xml"
+                location="https://mock-goal-state/zrdfepirv2sn6prdstr05a.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions.Edp_ExampleExtensionLinuxInTest_ussouth_manifest.xml"
                 state="enabled" autoUpgrade="true"
-                failoverlocation="https://rdfepirv2sn3prdstr01.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions.Edp_ExampleExtensionLinuxInTest_ussouth_manifest.xml"
+                failoverlocation="https://mock-goal-state/rdfepirv2sn3prdstr01.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions.Edp_ExampleExtensionLinuxInTest_ussouth_manifest.xml"
                 runAsStartupTask="false" isJson="true" useExactVersion="false"/>
     </Plugins>
     <PluginSettings>
diff --git a/tests/data/wire/ext_conf_no_public.xml b/tests/data/wire/ext_conf_no_public.xml
index d96b9f53..23ca595e 100644
--- a/tests/data/wire/ext_conf_no_public.xml
+++ b/tests/data/wire/ext_conf_no_public.xml
@@ -3,39 +3,39 @@
     <GAFamily>
       <Name>Win8</Name>
       <Uris>
-        <Uri>http://rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr05.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr06.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr07.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr08.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr09.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr10.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr11.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr12.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://zrdfepirv2hk2prdstr01.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr05.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr06.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr07.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr08.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr09.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr10.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr11.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr12.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/zrdfepirv2hk2prdstr01.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
       </Uris>
     </GAFamily>
     <GAFamily>
       <Name>Win7</Name>
       <Uris>
-        <Uri>http://rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr05.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr06.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr07.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr08.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr09.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr10.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr11.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr12.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://zrdfepirv2hk2prdstr01.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr05.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr06.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr07.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr08.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr09.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr10.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr11.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr12.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/zrdfepirv2hk2prdstr01.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
       </Uris>
     </GAFamily>
   </GAFamilies>
 </GuestAgentExtension>
 <Plugins>
-  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0" location="http://rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="false" failoverlocation="http://rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
+  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0" location="http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="false" failoverlocation="http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
 </Plugins>
 <PluginSettings>
   <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0">
diff --git a/tests/data/wire/ext_conf_no_settings.xml b/tests/data/wire/ext_conf_no_settings.xml
index 16972250..ed07fde5 100644
--- a/tests/data/wire/ext_conf_no_settings.xml
+++ b/tests/data/wire/ext_conf_no_settings.xml
@@ -3,39 +3,39 @@
     <GAFamily>
       <Name>Win8</Name>
       <Uris>
-        <Uri>http://rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr05.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr06.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr07.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr08.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr09.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr10.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr11.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr12.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
-        <Uri>http://zrdfepirv2hk2prdstr01.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr05.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr06.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr07.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr08.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr09.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr10.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr11.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr12.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/zrdfepirv2hk2prdstr01.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win8_asiaeast_manifest.xml</Uri>
       </Uris>
     </GAFamily>
     <GAFamily>
       <Name>Win7</Name>
       <Uris>
-        <Uri>http://rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr05.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr06.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr07.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr08.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr09.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr10.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr11.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://rdfepirv2hknprdstr12.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
-        <Uri>http://zrdfepirv2hk2prdstr01.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr05.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr06.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr07.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr08.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr09.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr10.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr11.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/rdfepirv2hknprdstr12.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
+        <Uri>http://mock-goal-state/zrdfepirv2hk2prdstr01.blob.core.windows.net/bfd5c281a7dc4e4b84381eb0b47e3aaf/Microsoft.WindowsAzure.GuestAgent_Win7_asiaeast_manifest.xml</Uri>
       </Uris>
     </GAFamily>
   </GAFamilies>
 </GuestAgentExtension>
 <Plugins>
-  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0" location="http://rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="false" failoverlocation="http://rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
+  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0" location="http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="false" failoverlocation="http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
 </Plugins>
 <StatusUploadBlob>https://test.blob.core.windows.net/vhds/test-cs12.test-cs12.test-cs12.status?sr=b&amp;sp=rw&amp;se=9999-01-01&amp;sk=key1&amp;sv=2014-02-14&amp;sig=hfRh7gzUE7sUtYwke78IOlZOrTRCYvkec4hGZ9zZzXo</StatusUploadBlob></Extensions>
 
diff --git a/tests/data/wire/ext_conf_sequencing.xml b/tests/data/wire/ext_conf_sequencing.xml
index bb215d0d..01c59e9b 100644
--- a/tests/data/wire/ext_conf_sequencing.xml
+++ b/tests/data/wire/ext_conf_sequencing.xml
@@ -3,20 +3,20 @@
     <GAFamily>
       <Name>Prod</Name>
       <Uris>
-          <Uri>http://manifest_of_ga.xml</Uri>
+          <Uri>http://mock-goal-state/manifest_of_ga.xml</Uri>
       </Uris>
     </GAFamily>
     <GAFamily>
       <Name>Test</Name>
       <Uris>
-          <Uri>http://manifest_of_ga.xml</Uri>
+          <Uri>http://mock-goal-state/manifest_of_ga.xml</Uri>
         </Uris>
     </GAFamily>
   </GAFamilies>
 </GuestAgentExtension>
 <Plugins>
-  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0" location="http://rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="false" failoverlocation="http://rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
-  <Plugin name="OSTCExtensions.OtherExampleHandlerLinux" version="1.0.0" location="http://rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_OtherExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="false" failoverlocation="http://rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_OtherExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
+  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0" location="http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="false" failoverlocation="http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
+  <Plugin name="OSTCExtensions.OtherExampleHandlerLinux" version="1.0.0" location="http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_OtherExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="false" failoverlocation="http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_OtherExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" />
 </Plugins>
 <PluginSettings>
   <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0">
diff --git a/tests/data/wire/ext_conf_upgradeguid.xml b/tests/data/wire/ext_conf_upgradeguid.xml
index e4c30daf..53994a2f 100644
--- a/tests/data/wire/ext_conf_upgradeguid.xml
+++ b/tests/data/wire/ext_conf_upgradeguid.xml
@@ -3,19 +3,19 @@
     <GAFamily>
       <Name>Prod</Name>
       <Uris>
-          <Uri>http://manifest_of_ga.xml</Uri>
+          <Uri>http://mock-goal-state/manifest_of_ga.xml</Uri>
       </Uris>
     </GAFamily>
     <GAFamily>
       <Name>Test</Name>
       <Uris>
-          <Uri>http://manifest_of_ga.xml</Uri>
+          <Uri>http://mock-goal-state/manifest_of_ga.xml</Uri>
         </Uris>
     </GAFamily>
   </GAFamilies>
 </GuestAgentExtension>
 <Plugins>
-  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0" location="http://rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="false" failoverlocation="http://rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" upgradeGuid="12345678-09AB-ABCD-CDEF-FE0987654321" />
+  <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0" location="http://mock-goal-state/rdfepirv2hknprdstr03.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" config="" state="enabled" autoUpgrade="false" failoverlocation="http://mock-goal-state/rdfepirv2hknprdstr04.blob.core.windows.net/b01058962be54ceca550a390fa5ff064/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml" runAsStartupTask="false" isJson="true" upgradeGuid="12345678-09AB-ABCD-CDEF-FE0987654321" />
 </Plugins>
 <PluginSettings>
   <Plugin name="OSTCExtensions.ExampleHandlerLinux" version="1.0.0">
diff --git a/tests/data/wire/ga_manifest.xml b/tests/data/wire/ga_manifest.xml
index fb7238a0..aae47b17 100644
--- a/tests/data/wire/ga_manifest.xml
+++ b/tests/data/wire/ga_manifest.xml
@@ -4,56 +4,56 @@
         <Plugin>
             <Version>1.0.0</Version>
             <Uris>
-                <Uri>http://foo.bar/zar/OSTCExtensions.WALinuxAgent__1.0.0</Uri>
+                <Uri>http://mock-goal-state/foo.bar/zar/OSTCExtensions.WALinuxAgent__1.0.0</Uri>
             </Uris>
         </Plugin>
         <Plugin>
             <Version>1.1.0</Version>
             <Uris>
-                <Uri>http://foo.bar/zar/OSTCExtensions.WALinuxAgent__1.1.0</Uri>
+                <Uri>http://mock-goal-state/foo.bar/zar/OSTCExtensions.WALinuxAgent__1.1.0</Uri>
             </Uris>
         </Plugin>
         <Plugin>
             <Version>1.1.1</Version>
             <Uris>
-                <Uri>http://foo.bar/zar/OSTCExtensions.WALinuxAgent__1.1.1</Uri>
+                <Uri>http://mock-goal-state/foo.bar/zar/OSTCExtensions.WALinuxAgent__1.1.1</Uri>
             </Uris>
         </Plugin>
         <Plugin>
             <Version>1.2.0</Version>
             <Uris>
-                <Uri>http://foo.bar/zar/OSTCExtensions.WALinuxAgent__1.2.0</Uri>
+                <Uri>http://mock-goal-state/foo.bar/zar/OSTCExtensions.WALinuxAgent__1.2.0</Uri>
             </Uris>
         </Plugin>
         <Plugin>
-            <Version>2.0.0</Version><Uris><Uri>http://host/OSTCExtensions.WALinuxAgent__2.0.0</Uri></Uris>
+            <Version>2.0.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.WALinuxAgent__2.0.0</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>2.1.0</Version><Uris><Uri>http://host/OSTCExtensions.WALinuxAgent__2.1.0</Uri></Uris>
+            <Version>2.1.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.WALinuxAgent__2.1.0</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>2.1.1</Version><Uris><Uri>http://host/OSTCExtensions.WALinuxAgent__2.1.1</Uri></Uris>
+            <Version>2.1.1</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.WALinuxAgent__2.1.1</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>2.2.0</Version><Uris><Uri>http://host/OSTCExtensions.WALinuxAgent__2.2.0</Uri></Uris>
+            <Version>2.2.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.WALinuxAgent__2.2.0</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>3.0</Version><Uris><Uri>http://host/OSTCExtensions.WALinuxAgent__3.0</Uri></Uris>
+            <Version>3.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.WALinuxAgent__3.0</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>3.1</Version><Uris><Uri>http://host/OSTCExtensions.WALinuxAgent__3.1</Uri></Uris>
+            <Version>3.1</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.WALinuxAgent__3.1</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>4.0.0.0</Version><Uris><Uri>http://host/OSTCExtensions.WALinuxAgent__3.0</Uri></Uris>
+            <Version>4.0.0.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.WALinuxAgent__3.0</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>4.0.0.1</Version><Uris><Uri>http://host/OSTCExtensions.WALinuxAgent__3.1</Uri></Uris>
+            <Version>4.0.0.1</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.WALinuxAgent__3.1</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>4.1.0.0</Version><Uris><Uri>http://host/OSTCExtensions.WALinuxAgent__3.1</Uri></Uris>
+            <Version>4.1.0.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.WALinuxAgent__3.1</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>99999.0.0.0</Version><Uris><Uri>http://host/OSTCExtensions.WALinuxAgent__99999.0.0.0</Uri></Uris>
+            <Version>99999.0.0.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.WALinuxAgent__99999.0.0.0</Uri></Uris>
         </Plugin>
     </Plugins>
 </PluginVersionManifest>
diff --git a/tests/data/wire/goal_state.xml b/tests/data/wire/goal_state.xml
index 0fd144f2..8c60c1cd 100644
--- a/tests/data/wire/goal_state.xml
+++ b/tests/data/wire/goal_state.xml
@@ -15,11 +15,11 @@
          <InstanceId>MachineRole_IN_0</InstanceId>
          <State>Started</State>
          <Configuration>
-         <HostingEnvironmentConfig>http://hostingenvuri/</HostingEnvironmentConfig>
-         <SharedConfig>http://sharedconfiguri/</SharedConfig>
-         <Certificates>http://certificatesuri/</Certificates>
-         <ExtensionsConfig>http://extensionsconfiguri/</ExtensionsConfig>
-         <FullConfig>http://fullconfiguri/</FullConfig>
+         <HostingEnvironmentConfig>http://168.63.129.16:80/hostingenvuri/</HostingEnvironmentConfig>
+         <SharedConfig>http://168.63.129.16:80/sharedconfiguri/</SharedConfig>
+         <Certificates>http://168.63.129.16:80/certificatesuri/</Certificates>
+         <ExtensionsConfig>http://168.63.129.16:80/extensionsconfiguri/</ExtensionsConfig>
+         <FullConfig>http://168.63.129.16:80/fullconfiguri/</FullConfig>
          <ConfigName>DummyRoleConfigName.xml</ConfigName>
          </Configuration>
        </RoleInstance>
diff --git a/tests/data/wire/goal_state_no_ext.xml b/tests/data/wire/goal_state_no_ext.xml
index 0b4f5669..1c267d9d 100644
--- a/tests/data/wire/goal_state_no_ext.xml
+++ b/tests/data/wire/goal_state_no_ext.xml
@@ -15,10 +15,10 @@
          <InstanceId>MachineRole_IN_0</InstanceId>
          <State>Started</State>
          <Configuration>
-         <HostingEnvironmentConfig>http://hostingenvuri/</HostingEnvironmentConfig>
-         <SharedConfig>http://sharedconfiguri/</SharedConfig>
-         <Certificates>http://certificatesuri/</Certificates>
-         <FullConfig>http://fullconfiguri/</FullConfig>
+         <HostingEnvironmentConfig>http://168.63.129.16:80/hostingenvuri/</HostingEnvironmentConfig>
+         <SharedConfig>http://168.63.129.16:80/sharedconfiguri/</SharedConfig>
+         <Certificates>http://168.63.129.16:80/certificatesuri/</Certificates>
+         <FullConfig>http://168.63.129.16:80/fullconfiguri/</FullConfig>
          </Configuration>
        </RoleInstance>
      </RoleInstanceList>
diff --git a/tests/data/wire/goal_state_remote_access.xml b/tests/data/wire/goal_state_remote_access.xml
index ba46ce2a..8d0a5cc8 100644
--- a/tests/data/wire/goal_state_remote_access.xml
+++ b/tests/data/wire/goal_state_remote_access.xml
@@ -10,18 +10,18 @@
    </Machine>
    <Container>
      <ContainerId>c6d5526c-5ac2-4200-b6e2-56f2b70c5ab2</ContainerId>
-     <RemoteAccessInfo>http://remoteaccessinfouri/
+     <RemoteAccessInfo>http://168.63.129.16:80/remoteaccessinfouri/
      </RemoteAccessInfo>
      <RoleInstanceList>
        <RoleInstance>
          <InstanceId>MachineRole_IN_0</InstanceId>
          <State>Started</State>
          <Configuration>
-         <HostingEnvironmentConfig>http://hostingenvuri/</HostingEnvironmentConfig>
-         <SharedConfig>http://sharedconfiguri/</SharedConfig>
-         <Certificates>http://certificatesuri/</Certificates>
-         <ExtensionsConfig>http://extensionsconfiguri/</ExtensionsConfig>
-         <FullConfig>http://fullconfiguri/</FullConfig>
+         <HostingEnvironmentConfig>http://168.63.129.16:80/hostingenvuri/</HostingEnvironmentConfig>
+         <SharedConfig>http://168.63.129.16:80/sharedconfiguri/</SharedConfig>
+         <Certificates>http://168.63.129.16:80/certificatesuri/</Certificates>
+         <ExtensionsConfig>http://168.63.129.16:80/extensionsconfiguri/</ExtensionsConfig>
+         <FullConfig>http://168.63.129.16:80/fullconfiguri/</FullConfig>
          <ConfigName>DummyRoleConfigName.xml</ConfigName>
          </Configuration>
        </RoleInstance>
diff --git a/tests/data/wire/manifest.xml b/tests/data/wire/manifest.xml
index eef8adee..dc32a50a 100644
--- a/tests/data/wire/manifest.xml
+++ b/tests/data/wire/manifest.xml
@@ -4,68 +4,68 @@
         <Plugin>
             <Version>1.0.0</Version>
             <Uris>
-                <Uri>http://foo.bar/zar/OSTCExtensions.ExampleHandlerLinux__1.0.0</Uri>
+                <Uri>http://mock-goal-state/foo.bar/zar/OSTCExtensions.ExampleHandlerLinux__1.0.0</Uri>
             </Uris>
         </Plugin>
         <Plugin>
             <Version>1.1.0</Version>
             <Uris>
-                <Uri>http://foo.bar/zar/OSTCExtensions.ExampleHandlerLinux__1.1.0</Uri>
+                <Uri>http://mock-goal-state/foo.bar/zar/OSTCExtensions.ExampleHandlerLinux__1.1.0</Uri>
             </Uris>
         </Plugin>
         <Plugin>
             <Version>1.1.1</Version>
             <Uris>
-                <Uri>http://foo.bar/zar/OSTCExtensions.ExampleHandlerLinux__1.1.1</Uri>
+                <Uri>http://mock-goal-state/foo.bar/zar/OSTCExtensions.ExampleHandlerLinux__1.1.1</Uri>
             </Uris>
         </Plugin>
         <Plugin>
             <Version>1.2.0</Version>
             <Uris>
-                <Uri>http://foo.bar/zar/OSTCExtensions.ExampleHandlerLinux__1.2.0</Uri>
+                <Uri>http://mock-goal-state/foo.bar/zar/OSTCExtensions.ExampleHandlerLinux__1.2.0</Uri>
             </Uris>
         </Plugin>
         <Plugin>
-            <Version>2.0.0</Version><Uris><Uri>http://host/OSTCExtensions.ExampleHandlerLinux__2.0.0</Uri></Uris>
+            <Version>2.0.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.ExampleHandlerLinux__2.0.0</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>2.1.0</Version><Uris><Uri>http://host/OSTCExtensions.ExampleHandlerLinux__2.1.0</Uri></Uris>
+            <Version>2.1.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.ExampleHandlerLinux__2.1.0</Uri></Uris>
             <DisallowMajorVersionUpgrade>True</DisallowMajorVersionUpgrade>
         </Plugin>
         <Plugin>
-            <Version>2.1.1</Version><Uris><Uri>http://host/OSTCExtensions.ExampleHandlerLinux__2.1.1</Uri></Uris>
+            <Version>2.1.1</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.ExampleHandlerLinux__2.1.1</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>2.2.0</Version><Uris><Uri>http://host/OSTCExtensions.ExampleHandlerLinux__2.2.0</Uri></Uris>
+            <Version>2.2.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.ExampleHandlerLinux__2.2.0</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>3.0</Version><Uris><Uri>http://host/OSTCExtensions.ExampleHandlerLinux__3.0</Uri></Uris>
+            <Version>3.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.ExampleHandlerLinux__3.0</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>3.1</Version><Uris><Uri>http://host/OSTCExtensions.ExampleHandlerLinux__3.1</Uri></Uris>
+            <Version>3.1</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.ExampleHandlerLinux__3.1</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>4.0.0.0</Version><Uris><Uri>http://host/OSTCExtensions.ExampleHandlerLinux__3.0</Uri></Uris>
+            <Version>4.0.0.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.ExampleHandlerLinux__3.0</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>4.0.0.1</Version><Uris><Uri>http://host/OSTCExtensions.ExampleHandlerLinux__3.1</Uri></Uris>
+            <Version>4.0.0.1</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.ExampleHandlerLinux__3.1</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>4.1.0.0</Version><Uris><Uri>http://host/OSTCExtensions.ExampleHandlerLinux__3.1</Uri></Uris>
+            <Version>4.1.0.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.ExampleHandlerLinux__3.1</Uri></Uris>
         </Plugin>
     </Plugins>
     <InternalPlugins>
         <Plugin>
             <Version>1.3.0</Version>
             <Uris>
-                <Uri>http://foo.bar/zar/OSTCExtensions.ExampleHandlerLinux__1.3.0</Uri>
+                <Uri>http://mock-goal-state/foo.bar/zar/OSTCExtensions.ExampleHandlerLinux__1.3.0</Uri>
             </Uris>
         </Plugin>
         <Plugin>
-            <Version>2.3.0</Version><Uris><Uri>http://host/OSTCExtensions.ExampleHandlerLinux__2.3.0</Uri></Uris>
+            <Version>2.3.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.ExampleHandlerLinux__2.3.0</Uri></Uris>
         </Plugin>
         <Plugin>
-            <Version>2.4.0</Version><Uris><Uri>http://host/OSTCExtensions.ExampleHandlerLinux__2.3.0</Uri></Uris>
+            <Version>2.4.0</Version><Uris><Uri>http://mock-goal-state/host/OSTCExtensions.ExampleHandlerLinux__2.3.0</Uri></Uris>
         </Plugin>
     </InternalPlugins>
 </PluginVersionManifest>
diff --git a/tests/data/wire/manifest_deletion.xml b/tests/data/wire/manifest_deletion.xml
index b79bad0f..ce655b60 100644
--- a/tests/data/wire/manifest_deletion.xml
+++ b/tests/data/wire/manifest_deletion.xml
@@ -4,7 +4,7 @@
         <Plugin>
             <Version>1.0.0</Version>
             <Uris>
-                <Uri>http://foo.bar/zar/OSTCExtensions.ExampleHandlerLinux__1.0.0</Uri>
+                <Uri>http://mock-goal-state/foo.bar/zar/OSTCExtensions.ExampleHandlerLinux__1.0.0</Uri>
             </Uris>
         </Plugin>
     </Plugins>
diff --git a/tests/ga/test_extension.py b/tests/ga/test_extension.py
index 58e636a5..06823b49 100644
--- a/tests/ga/test_extension.py
+++ b/tests/ga/test_extension.py
@@ -33,6 +33,7 @@ from azurelinuxagent.common.exception import ResourceGoneError
 from azurelinuxagent.common.protocol.restapi import Extension, ExtHandlerProperties
 from azurelinuxagent.ga.exthandlers import *
 from azurelinuxagent.common.protocol.wire import WireProtocol, InVMArtifactsProfile
+from azurelinuxagent.common.utils.restutil import KNOWN_WIRESERVER_IP
 
 # Mocking the original sleep to reduce test execution time
 SLEEP = time.sleep
@@ -348,7 +349,7 @@ class TestExtension(ExtensionTestCase):
         mock_http_get.side_effect = test_data.mock_http_get
         MockCryptUtil.side_effect = test_data.mock_crypt_util
 
-        protocol = WireProtocol("foo.bar")
+        protocol = WireProtocol(KNOWN_WIRESERVER_IP)
         protocol.detect()
         protocol.report_ext_status = MagicMock()
         protocol.report_vm_status = MagicMock()
@@ -1259,7 +1260,7 @@ class TestExtension(ExtensionTestCase):
 
         _, protocol = self._create_mock(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE), *args)
         version_uri = Mock()
-        version_uri.uri = 'http://some/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml'
+        version_uri.uri = 'http://mock-goal-state/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml'
 
         for (installed_version, config_version, expected_version) in cases:
             ext_handler = Mock()
@@ -1988,7 +1989,7 @@ class TestExtensionSequencing(AgentTestCase):
         mock_http_get.side_effect = test_data.mock_http_get
         MockCryptUtil.side_effect = test_data.mock_crypt_util
 
-        protocol = WireProtocol("foo.bar")
+        protocol = WireProtocol(KNOWN_WIRESERVER_IP)
         protocol.detect()
         protocol.report_ext_status = MagicMock()
         protocol.report_vm_status = MagicMock()
@@ -2223,7 +2224,7 @@ class TestExtensionWithCGroupsEnabled(AgentTestCase):
         mock_http_get.side_effect = test_data.mock_http_get
         mock_crypt_util.side_effect = test_data.mock_crypt_util
 
-        protocol = WireProtocol("foo.bar")
+        protocol = WireProtocol(KNOWN_WIRESERVER_IP)
         protocol.detect()
         protocol.report_ext_status = MagicMock()
         protocol.report_vm_status = MagicMock()
diff --git a/tests/ga/test_monitor.py b/tests/ga/test_monitor.py
index ad44934b..c65a75de 100644
--- a/tests/ga/test_monitor.py
+++ b/tests/ga/test_monitor.py
@@ -109,7 +109,7 @@ def get_event_message(duration, evt_type, is_internal, is_success, message, name
 
 @patch('azurelinuxagent.common.event.EventLogger.add_event')
 @patch('azurelinuxagent.common.osutil.get_osutil')
-@patch('azurelinuxagent.common.protocol.get_protocol_util')
+@patch('azurelinuxagent.common.protocol.util.get_protocol_util')
 @patch('azurelinuxagent.common.protocol.util.ProtocolUtil.get_protocol')
 @patch("azurelinuxagent.common.protocol.healthservice.HealthService._report")
 @patch("azurelinuxagent.common.utils.restutil.http_get")
@@ -567,7 +567,7 @@ class TestEventMonitoring(AgentTestCase):
         mock_http_get.side_effect = test_data.mock_http_get
         MockCryptUtil.side_effect = test_data.mock_crypt_util
 
-        protocol = WireProtocol("foo.bar")
+        protocol = WireProtocol(restutil.KNOWN_WIRESERVER_IP)
         protocol.detect()
         protocol.report_ext_status = MagicMock()
         protocol.report_vm_status = MagicMock()
@@ -827,7 +827,7 @@ class TestEventMonitoring(AgentTestCase):
                 monitor_handler.collect_and_send_events()
                 self.assertEqual(1, mock_error.call_count)
                 self.assertEqual("[ProtocolError] [Wireserver Exception] [ProtocolError] [Wireserver Failed] "
-                                 "URI http://foo.bar/machine?comp=telemetrydata  [HTTP Failed] Status Code 503",
+                                 "URI http://{0}/machine?comp=telemetrydata  [HTTP Failed] Status Code 503".format(protocol.get_endpoint()),
                                  mock_error.call_args[0][1])
                 self.assertEqual(0, len(os.listdir(self.event_logger.event_dir)))
 
@@ -882,7 +882,7 @@ class TestEventMonitoring(AgentTestCase):
 
 
 @patch('azurelinuxagent.common.osutil.get_osutil')
-@patch('azurelinuxagent.common.protocol.get_protocol_util')
+@patch('azurelinuxagent.common.protocol.util.get_protocol_util')
 @patch("azurelinuxagent.common.protocol.healthservice.HealthService._report")
 @patch("azurelinuxagent.common.utils.restutil.http_get")
 class TestExtensionMetricsDataTelemetry(AgentTestCase):
diff --git a/tests/ga/test_remoteaccess.py b/tests/ga/test_remoteaccess.py
index 05ee25f6..7d8544ef 100644
--- a/tests/ga/test_remoteaccess.py
+++ b/tests/ga/test_remoteaccess.py
@@ -16,8 +16,9 @@
 #
 import xml
 
-from tests.tools import AgentTestCase, load_data, patch
-from azurelinuxagent.common.protocol.wire import *
+from azurelinuxagent.common.protocol.goal_state import GoalState, RemoteAccess
+from tests.tools import AgentTestCase, load_data, patch, Mock
+from tests.protocol import mockwiredata, mock_wire_protocol
 
 
 class TestRemoteAccess(AgentTestCase):
@@ -31,12 +32,9 @@ class TestRemoteAccess(AgentTestCase):
         self.assertEquals("encryptedPasswordString", remote_access.user_list.users[0].encrypted_password, "Encrypted password does not match.")
         self.assertEquals("2019-01-01", remote_access.user_list.users[0].expiration, "Expiration does not match.")
 
-    @patch('azurelinuxagent.common.protocol.wire.WireClient.get_goal_state',
-    return_value=GoalState(load_data('wire/goal_state.xml')))
-    def test_update_remote_access_conf_no_remote_access(self, _):
-        protocol = WireProtocol('12.34.56.78')
-        goal_state = protocol.client.get_goal_state()
-        protocol.client.update_remote_access_conf(goal_state)
+    def test_goal_state_with_no_remote_access(self):
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE) as protocol:
+            self.assertIsNone(protocol.client.get_remote_access())
 
     def test_parse_two_remote_access_accounts(self):
         data_str = load_data('wire/remote_access_two_accounts.xml')
@@ -56,7 +54,7 @@ class TestRemoteAccess(AgentTestCase):
         remote_access = RemoteAccess(data_str)
         self.assertNotEquals(None, remote_access)
         self.assertEquals(10, len(remote_access.user_list.users), "User count does not match.")
-    
+
     def test_parse_duplicate_remote_access_accounts(self):
         data_str = load_data('wire/remote_access_duplicate_accounts.xml')
         remote_access = RemoteAccess(data_str)
@@ -75,19 +73,12 @@ class TestRemoteAccess(AgentTestCase):
         self.assertNotEquals(None, remote_access)
         self.assertEquals(0, len(remote_access.user_list.users), "User count does not match.")
 
-    @patch('azurelinuxagent.common.protocol.wire.WireClient.get_goal_state',
-    return_value=GoalState(load_data('wire/goal_state_remote_access.xml')))
-    @patch('azurelinuxagent.common.protocol.wire.WireClient.fetch_config',
-    return_value=load_data('wire/remote_access_single_account.xml'))
-    @patch('azurelinuxagent.common.protocol.wire.WireClient.get_header_for_cert')
-    def test_update_remote_access_conf_remote_access(self, _1, _2, _3):
-        protocol = WireProtocol('12.34.56.78')
-        goal_state = protocol.client.get_goal_state()
-        protocol.client.update_remote_access_conf(goal_state)
-        self.assertNotEquals(None, protocol.client.get_remote_access())
-        self.assertEquals(1, len(protocol.client.get_remote_access().user_list.users))
-        self.assertEquals('testAccount', protocol.client.get_remote_access().user_list.users[0].name)
-        self.assertEquals('encryptedPasswordString', protocol.client.get_remote_access().user_list.users[0].encrypted_password)
+    def test_update_remote_access_conf_remote_access(self):
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE_REMOTE_ACCESS) as protocol:
+            self.assertIsNotNone(protocol.client.get_remote_access())
+            self.assertEquals(1, len(protocol.client.get_remote_access().user_list.users))
+            self.assertEquals('testAccount', protocol.client.get_remote_access().user_list.users[0].name)
+            self.assertEquals('encryptedPasswordString', protocol.client.get_remote_access().user_list.users[0].encrypted_password)
 
     def test_parse_bad_remote_access_data(self):
         data = "foobar"
diff --git a/tests/ga/test_remoteaccess_handler.py b/tests/ga/test_remoteaccess_handler.py
index 7d569362..b67b87d6 100644
--- a/tests/ga/test_remoteaccess_handler.py
+++ b/tests/ga/test_remoteaccess_handler.py
@@ -21,6 +21,7 @@ from mock import Mock, MagicMock
 from azurelinuxagent.common.protocol.util import ProtocolUtil
 
 from azurelinuxagent.common.exception import RemoteAccessError
+from azurelinuxagent.common.protocol.goal_state import RemoteAccess
 from azurelinuxagent.common.protocol.wire import *
 from azurelinuxagent.ga.remoteaccess import RemoteAccessHandler
 from tests.common.osutil.mock_osutil import MockOSUtil
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index 002742f3..4a8d8ce4 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -10,6 +10,7 @@ from threading import currentThread
 
 from azurelinuxagent.ga.env import EnvHandler
 
+from azurelinuxagent.common.protocol.goal_state import ExtensionsConfig
 from azurelinuxagent.common.protocol.hostplugin import *
 from azurelinuxagent.common.protocol.util import ProtocolUtil
 from azurelinuxagent.common.protocol.wire import *
diff --git a/tests/pa/test_deprovision.py b/tests/pa/test_deprovision.py
index 591a2c67..3d484954 100644
--- a/tests/pa/test_deprovision.py
+++ b/tests/pa/test_deprovision.py
@@ -29,7 +29,7 @@ from tests.tools import AgentTestCase, distros, Mock, patch
 class TestDeprovision(AgentTestCase):
     @patch('signal.signal')
     @patch('azurelinuxagent.common.osutil.get_osutil')
-    @patch('azurelinuxagent.common.protocol.get_protocol_util')
+    @patch('azurelinuxagent.common.protocol.util.get_protocol_util')
     @patch('azurelinuxagent.pa.deprovision.default.read_input')
     def test_confirmation(self,
             mock_read, mock_protocol, mock_util, mock_signal):
diff --git a/tests/pa/test_provision.py b/tests/pa/test_provision.py
index 54175601..6bb98577 100644
--- a/tests/pa/test_provision.py
+++ b/tests/pa/test_provision.py
@@ -23,7 +23,7 @@ import unittest
 import azurelinuxagent.common.conf as conf
 from azurelinuxagent.common.exception import ProvisionError
 from azurelinuxagent.common.osutil.default import DefaultOSUtil
-from azurelinuxagent.common.protocol import OVF_FILE_NAME
+from azurelinuxagent.common.protocol.util import OVF_FILE_NAME
 from azurelinuxagent.pa.provision import get_provision_handler
 from azurelinuxagent.pa.provision.cloudinit import CloudInitProvisionHandler
 from azurelinuxagent.pa.provision.default import ProvisionHandler
diff --git a/tests/protocol/mock_wire_protocol.py b/tests/protocol/mock_wire_protocol.py
new file mode 100644
index 00000000..cafd7386
--- /dev/null
+++ b/tests/protocol/mock_wire_protocol.py
@@ -0,0 +1,84 @@
+# Copyright 2020 Microsoft Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Requires Python 2.6+ and Openssl 1.0+
+#
+
+import contextlib
+import re
+from azurelinuxagent.common.protocol.wire import WireProtocol
+from azurelinuxagent.common.utils.restutil import KNOWN_WIRESERVER_IP, http_get
+from tests.tools import patch
+from tests.protocol.mockwiredata import WireProtocolData
+
+@contextlib.contextmanager
+def create(mock_wire_data_file):
+    """
+    Creates a mock WireProtocol object that will return the data specified by 'mock_wire_data_file' (which must
+    follow the structure of the data files defined in tests/protocol/mockwiredata.py).
+
+    NOTE: This function creates mocks for azurelinuxagent.common.utils.restutil.http_get and
+          azurelinuxagent.common.protocol.wire.CryptUtil. These mocks can be stopped using the
+          methods stop_mock_http_get() and stop_mock_crypt_util().
+
+    The return value is an instance of WireProtocol augmented with these properties/methods:
+
+        * mock_wire_data - the WireProtocolData constructed from the mock_wire_data_file parameter.
+        * stop_mock_http_get() - stops the mock for restutil.http_get
+        * stop_mock_crypt_util() - stops the mock for CrypUtil
+    """
+    def stop_mock_http_get():
+        if stop_mock_http_get.mock is not None:
+            stop_mock_http_get.mock.stop()
+            stop_mock_http_get.mock = None
+    stop_mock_http_get.mock = None
+
+    def stop_mock_crypt_util():
+        if stop_mock_crypt_util.mock is not None:
+            stop_mock_crypt_util.mock.stop()
+            stop_mock_crypt_util.mock = None
+    stop_mock_crypt_util.mock = None
+
+    protocol = WireProtocol(KNOWN_WIRESERVER_IP)
+    protocol.mock_wire_data = WireProtocolData(mock_wire_data_file)
+    protocol.stop_mock_http_get = stop_mock_http_get
+    protocol.stop_mock_crypt_util = stop_mock_crypt_util
+
+    try:
+        # To minimize the impact of mocking restutil.http_get we only use the mock data for requests
+        # to the wireserver or requests starting with "mock-goal-state"
+        original_http_get = http_get
+
+        mock_data_re = re.compile(r'https?://(mock-goal-state|{0}).*'.format(KNOWN_WIRESERVER_IP.replace(r'.', r'\.')), re.IGNORECASE)
+
+        def mock_http_get(url, *args, **kwargs):
+            if mock_data_re.match(url) is None:
+                return original_http_get(url, *args, **kwargs)
+            return protocol.mock_wire_data.mock_http_get(url, *args, **kwargs)
+
+        p = patch("azurelinuxagent.common.utils.restutil.http_get", side_effect=mock_http_get)
+        p.start()
+        stop_mock_http_get.mock = p
+
+        p = patch("azurelinuxagent.common.protocol.wire.CryptUtil", side_effect=protocol.mock_wire_data.mock_crypt_util)
+        p.start()
+        stop_mock_crypt_util.mock = p
+
+        protocol.detect()
+
+        yield protocol
+
+    finally:
+        protocol.stop_mock_crypt_util()
+        protocol.stop_mock_http_get()
diff --git a/tests/protocol/mockwiredata.py b/tests/protocol/mockwiredata.py
index fdfc0fdb..5c8833ee 100644
--- a/tests/protocol/mockwiredata.py
+++ b/tests/protocol/mockwiredata.py
@@ -33,7 +33,8 @@ DATA_FILE = {
         "ga_manifest": "wire/ga_manifest.xml",
         "trans_prv": "wire/trans_prv",
         "trans_cert": "wire/trans_cert",
-        "test_ext": "ext/sample_ext-1.3.0.zip"
+        "test_ext": "ext/sample_ext-1.3.0.zip",
+        "remote_access": None
 }
 
 DATA_FILE_NO_EXT = DATA_FILE.copy()
@@ -75,6 +76,10 @@ DATA_FILE_NO_CERT_FORMAT["certs"] = "wire/certs_no_format_specified.xml"
 DATA_FILE_CERT_FORMAT_NOT_PFX = DATA_FILE.copy()
 DATA_FILE_CERT_FORMAT_NOT_PFX["certs"] = "wire/certs_format_not_pfx.xml"
 
+DATA_FILE_REMOTE_ACCESS = DATA_FILE.copy()
+DATA_FILE_REMOTE_ACCESS["goal_state"] = "wire/goal_state_remote_access.xml"
+DATA_FILE_REMOTE_ACCESS["remote_access"] = "wire/remote_access_single_account.xml"
+
 class WireProtocolData(object):
     def __init__(self, data_files=DATA_FILE):
         self.emulate_stale_goal_state = False
@@ -86,6 +91,7 @@ class WireProtocolData(object):
             "sharedconfiguri": 0,
             "certificatesuri": 0,
             "extensionsconfiguri": 0,
+            "remoteaccessinfouri": 0,
             "extensionArtifact": 0,
             "manifest.xml": 0,
             "manifest_of_ga.xml": 0,
@@ -103,6 +109,7 @@ class WireProtocolData(object):
         self.trans_prv = None
         self.trans_cert = None
         self.ext = None
+        self.remote_access = None
 
         self.reload()
 
@@ -118,6 +125,9 @@ class WireProtocolData(object):
         self.trans_prv = load_data(self.data_files.get("trans_prv"))
         self.trans_cert = load_data(self.data_files.get("trans_cert"))
         self.ext = load_bin_data(self.data_files.get("test_ext"))
+        remote_access_data_file = self.data_files.get("remote_access")
+        if remote_access_data_file is not None:
+            self.remote_access = load_data(remote_access_data_file)
 
     def mock_http_get(self, url, *args, **kwargs):
         content = None
@@ -149,6 +159,9 @@ class WireProtocolData(object):
         elif "extensionsconfiguri" in url:
             content = self.ext_conf
             self.call_counts["extensionsconfiguri"] += 1
+        elif "remoteaccessinfouri" in url:
+            content = self.remote_access
+            self.call_counts["remoteaccessinfouri"] += 1
 
         else:
             # A stale GoalState results in a 400 from the HostPlugin
diff --git a/tests/protocol/test_hostplugin.py b/tests/protocol/test_hostplugin.py
index f0fac933..58c7e25d 100644
--- a/tests/protocol/test_hostplugin.py
+++ b/tests/protocol/test_hostplugin.py
@@ -20,6 +20,7 @@ import json
 import sys
 import datetime
 import unittest
+import contextlib
 
 import azurelinuxagent.common.protocol.restapi as restapi
 import azurelinuxagent.common.protocol.wire as wire
@@ -30,7 +31,8 @@ from azurelinuxagent.common.exception import HttpError, ResourceGoneError
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.protocol.hostplugin import API_VERSION
 from azurelinuxagent.common.utils import restutil
-from tests.protocol.mockwiredata import WireProtocolData, DATA_FILE
+from tests.protocol import mock_wire_protocol
+from tests.protocol.mockwiredata import WireProtocolData, DATA_FILE, DATA_FILE_NO_EXT
 from tests.protocol.test_wire import MockResponse
 from tests.tools import AgentTestCase, PY_VERSION_MAJOR, Mock, patch
 
@@ -63,12 +65,13 @@ if PY_VERSION_MAJOR > 2:
 class TestHostPlugin(AgentTestCase):
 
     def _init_host(self):
-        test_goal_state = wire.GoalState(WireProtocolData(DATA_FILE).goal_state)
-        host_plugin = wire.HostPluginProtocol(wireserver_url,
-                                              test_goal_state.container_id,
-                                              test_goal_state.role_config_name)
-        self.assertTrue(host_plugin.health_service is not None)
-        return host_plugin
+        with mock_wire_protocol.create(DATA_FILE) as protocol:
+            test_goal_state = protocol.client.get_goal_state()
+            host_plugin = wire.HostPluginProtocol(wireserver_url,
+                                                  test_goal_state.container_id,
+                                                  test_goal_state.role_config_name)
+            self.assertTrue(host_plugin.health_service is not None)
+            return host_plugin
 
     def _init_status_blob(self):
         wire_protocol_client = wire.WireProtocol(wireserver_url).client
@@ -149,20 +152,23 @@ class TestHostPlugin(AgentTestCase):
         self.assertEqual(headers['x-ms-containerid'], goal_state.container_id)
         self.assertEqual(headers['x-ms-host-config-name'], goal_state.role_config_name)
 
-    def _get_test_wire_protocol(self):
-        test_goal_state = wire.GoalState(WireProtocolData(DATA_FILE).goal_state)
-        status = restapi.VMStatus(status="Ready", message="Guest Agent is running")
+    @staticmethod
+    @contextlib.contextmanager
+    def create_mock_protocol():
+        with mock_wire_protocol.create(DATA_FILE_NO_EXT) as protocol:
+            # These tests use mock wire data that dont have any extensions (extension config will be empty).
+            # Populate the upload blob and set an initial empty status before returning the protocol.
+            ext_conf = protocol.client._goal_state.ext_conf
+            ext_conf.status_upload_blob = sas_url
+            ext_conf.status_upload_blob_type = page_blob_type
 
-        wire_protocol = wire.WireProtocol(wireserver_url)
-        wire_protocol_client = wire_protocol.client
-        wire_protocol_client.get_goal_state = Mock(return_value=test_goal_state)
-        ext_conf = wire.ExtensionsConfig(None)
-        ext_conf.status_upload_blob = sas_url
-        ext_conf.status_upload_blob_type = page_blob_type
-        wire_protocol_client._set_ext_conf(ext_conf)
-        wire_protocol_client.status_blob.set_vm_status(status)
+            status = restapi.VMStatus(status="Ready", message="Guest Agent is running")
+            protocol.client.status_blob.set_vm_status(status)
 
-        return wire_protocol
+            # Also, they mock WireClient.update_goal_state() to verify how it is called
+            protocol.client.update_goal_state = Mock()
+
+            yield protocol
 
     @patch("azurelinuxagent.common.protocol.healthservice.HealthService.report_host_plugin_versions")
     @patch("azurelinuxagent.ga.update.restutil.http_get")
@@ -237,36 +243,33 @@ class TestHostPlugin(AgentTestCase):
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.ensure_initialized", return_value=True)
     @patch("azurelinuxagent.common.protocol.wire.StatusBlob.upload", return_value=False)
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol._put_page_blob_status")
-    @patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state")
-    def test_default_channel(self, patch_update, patch_put, patch_upload, _):
+    def test_default_channel(self, patch_put, patch_upload, _):
         """
         Status now defaults to HostPlugin. Validate that any errors on the public
         channel are ignored.  Validate that the default channel is never changed
         as part of status upload.
         """
+        with self.create_mock_protocol() as wire_protocol:
+            wire_protocol.update_goal_state()
 
-        wire_protocol = self._get_test_wire_protocol()
-
-        wire_protocol.update_goal_state()
-
-        # act
-        wire_protocol.client.upload_status_blob()
+            # act
+            wire_protocol.client.upload_status_blob()
 
-        # assert direct route is not called
-        self.assertEqual(0, patch_upload.call_count, "Direct channel was used")
+            # assert direct route is not called
+            self.assertEqual(0, patch_upload.call_count, "Direct channel was used")
 
-        # assert host plugin route is called
-        self.assertEqual(1, patch_put.call_count, "Host plugin was not used")
+            # assert host plugin route is called
+            self.assertEqual(1, patch_put.call_count, "Host plugin was not used")
 
-        # assert update goal state is only called once, non-forced
-        self.assertEqual(1, patch_update.call_count, "Unexpected call count")
-        self.assertEqual(0, len(patch_update.call_args[1]), "Unexpected parameters")
+            # assert update goal state is only called once, non-forced
+            self.assertEqual(1, wire_protocol.client.update_goal_state.call_count, "Unexpected call count")
+            self.assertEqual(0, len(wire_protocol.client.update_goal_state.call_args[1]), "Unexpected parameters")
 
-        # ensure the correct url is used
-        self.assertEqual(sas_url, patch_put.call_args[0][0])
+            # ensure the correct url is used
+            self.assertEqual(sas_url, patch_put.call_args[0][0])
 
-        # ensure host plugin is not set as default
-        self.assertFalse(wire.HostPluginProtocol.is_default_channel())
+            # ensure host plugin is not set as default
+            self.assertFalse(wire.HostPluginProtocol.is_default_channel())
 
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.ensure_initialized",
            return_value=True)
@@ -274,33 +277,31 @@ class TestHostPlugin(AgentTestCase):
            return_value=True)
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol._put_page_blob_status",
            side_effect=HttpError("503"))
-    @patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state")
-    def test_fallback_channel_503(self, patch_update, patch_put, patch_upload, _):
+    def test_fallback_channel_503(self, patch_put, patch_upload, _):
         """
         When host plugin returns a 503, we should fall back to the direct channel
         """
-        wire_protocol = self._get_test_wire_protocol()
+        with self.create_mock_protocol() as wire_protocol:
+            wire_protocol.update_goal_state()
 
-        wire_protocol.update_goal_state()
+            # act
+            wire_protocol.client.upload_status_blob()
 
-        # act
-        wire_protocol.client.upload_status_blob()
+            # assert direct route is called
+            self.assertEqual(1, patch_upload.call_count, "Direct channel was not used")
 
-        # assert direct route is called
-        self.assertEqual(1, patch_upload.call_count, "Direct channel was not used")
+            # assert host plugin route is called
+            self.assertEqual(1, patch_put.call_count, "Host plugin was not used")
 
-        # assert host plugin route is called
-        self.assertEqual(1, patch_put.call_count, "Host plugin was not used")
+            # assert update goal state is only called once, non-forced
+            self.assertEqual(1, wire_protocol.client.update_goal_state.call_count, "Update goal state unexpected call count")
+            self.assertEqual(0, len(wire_protocol.client.update_goal_state.call_args[1]), "Update goal state unexpected call count")
 
-        # assert update goal state is only called once, non-forced
-        self.assertEqual(1, patch_update.call_count, "Update goal state unexpected call count")
-        self.assertEqual(0, len(patch_update.call_args[1]), "Update goal state unexpected call count")
+            # ensure the correct url is used
+            self.assertEqual(sas_url, patch_put.call_args[0][0])
 
-        # ensure the correct url is used
-        self.assertEqual(sas_url, patch_put.call_args[0][0])
-
-        # ensure host plugin is not set as default
-        self.assertFalse(wire.HostPluginProtocol.is_default_channel())
+            # ensure host plugin is not set as default
+            self.assertFalse(wire.HostPluginProtocol.is_default_channel())
 
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.ensure_initialized",
            return_value=True)
@@ -308,35 +309,33 @@ class TestHostPlugin(AgentTestCase):
            return_value=True)
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol._put_page_blob_status",
            side_effect=ResourceGoneError("410"))
-    @patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state")
     @patch("azurelinuxagent.common.protocol.wire.WireClient.update_host_plugin_from_goal_state")
-    def test_fallback_channel_410(self, patch_refresh_host_plugin, patch_update, patch_put, patch_upload, _):
+    def test_fallback_channel_410(self, patch_refresh_host_plugin, patch_put, patch_upload, _):
         """
         When host plugin returns a 410, we should force the goal state update and return
         """
-        wire_protocol = self._get_test_wire_protocol()
-
-        wire_protocol.update_goal_state()
+        with self.create_mock_protocol() as wire_protocol:
+            wire_protocol.update_goal_state()
 
-        # act
-        wire_protocol.client.upload_status_blob()
+            # act
+            wire_protocol.client.upload_status_blob()
 
-        # assert direct route is not called
-        self.assertEqual(0, patch_upload.call_count, "Direct channel was used")
+            # assert direct route is not called
+            self.assertEqual(0, patch_upload.call_count, "Direct channel was used")
 
-        # assert host plugin route is called
-        self.assertEqual(1, patch_put.call_count, "Host plugin was not used")
+            # assert host plugin route is called
+            self.assertEqual(1, patch_put.call_count, "Host plugin was not used")
 
-        # assert update goal state is called with no arguments (forced=False), then update_host_plugin_from_goal_state is called
-        self.assertEqual(1, patch_update.call_count, "Update goal state unexpected call count")
-        self.assertEqual(0, len(patch_update.call_args[1]), "Update goal state unexpected argument count")
-        self.assertEqual(1, patch_refresh_host_plugin.call_count, "Refresh host plugin unexpected call count")
+            # assert update goal state is called with no arguments (forced=False), then update_host_plugin_from_goal_state is called
+            self.assertEqual(1, wire_protocol.client.update_goal_state.call_count, "Update goal state unexpected call count")
+            self.assertEqual(0, len(wire_protocol.client.update_goal_state.call_args[1]), "Update goal state unexpected argument count")
+            self.assertEqual(1, patch_refresh_host_plugin.call_count, "Refresh host plugin unexpected call count")
 
-        # ensure the correct url is used
-        self.assertEqual(sas_url, patch_put.call_args[0][0])
+            # ensure the correct url is used
+            self.assertEqual(sas_url, patch_put.call_args[0][0])
 
-        # ensure host plugin is not set as default
-        self.assertFalse(wire.HostPluginProtocol.is_default_channel())
+            # ensure host plugin is not set as default
+            self.assertFalse(wire.HostPluginProtocol.is_default_channel())
 
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.ensure_initialized",
            return_value=True)
@@ -344,248 +343,226 @@ class TestHostPlugin(AgentTestCase):
            return_value=False)
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol._put_page_blob_status",
            side_effect=HttpError("500"))
-    @patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state")
-    def test_fallback_channel_failure(self, patch_update, patch_put, patch_upload, _):
+    def test_fallback_channel_failure(self, patch_put, patch_upload, _):
         """
         When host plugin returns a 500, and direct fails, we should raise a ProtocolError
         """
-        wire_protocol = self._get_test_wire_protocol()
+        with self.create_mock_protocol() as wire_protocol:
+            wire_protocol.update_goal_state()
 
-        wire_protocol.update_goal_state()
+            # act
+            self.assertRaises(wire.ProtocolError, wire_protocol.client.upload_status_blob)
 
-        # act
-        self.assertRaises(wire.ProtocolError, wire_protocol.client.upload_status_blob)
+            # assert direct route is not called
+            self.assertEqual(1, patch_upload.call_count, "Direct channel was not used")
 
-        # assert direct route is not called
-        self.assertEqual(1, patch_upload.call_count, "Direct channel was not used")
+            # assert host plugin route is called
+            self.assertEqual(1, patch_put.call_count, "Host plugin was not used")
 
-        # assert host plugin route is called
-        self.assertEqual(1, patch_put.call_count, "Host plugin was not used")
+            # assert update goal state is called twice, forced=True on the second
+            self.assertEqual(1, wire_protocol.client.update_goal_state.call_count, "Update goal state unexpected call count")
+            self.assertEqual(0, len(wire_protocol.client.update_goal_state.call_args[1]), "Update goal state unexpected call count")
 
-        # assert update goal state is called twice, forced=True on the second
-        self.assertEqual(1, patch_update.call_count, "Update goal state unexpected call count")
-        self.assertEqual(0, len(patch_update.call_args[1]), "Update goal state unexpected call count")
+            # ensure the correct url is used
+            self.assertEqual(sas_url, patch_put.call_args[0][0])
 
-        # ensure the correct url is used
-        self.assertEqual(sas_url, patch_put.call_args[0][0])
+            # ensure host plugin is not set as default
+            self.assertFalse(wire.HostPluginProtocol.is_default_channel())
 
-        # ensure host plugin is not set as default
-        self.assertFalse(wire.HostPluginProtocol.is_default_channel())
-
-    @patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state")
     @patch("azurelinuxagent.common.event.add_event")
-    def test_put_status_error_reporting(self, patch_add_event, _):
+    def test_put_status_error_reporting(self, patch_add_event):
         """
         Validate the telemetry when uploading status fails
         """
         wire.HostPluginProtocol.set_default_channel(False)
         with patch.object(wire.StatusBlob, "upload", return_value=False):
-
-            wire_protocol = self._get_test_wire_protocol()
-            wire_protocol_client = wire_protocol.client
-
-            put_error = wire.HttpError("put status http error")
-            with patch.object(restutil, "http_put", side_effect=put_error):
-                with patch.object(wire.HostPluginProtocol,
-                                  "ensure_initialized", return_value=True):
-                    self.assertRaises(wire.ProtocolError, wire_protocol_client.upload_status_blob)
-
-                    # The agent tries to upload via HostPlugin and that fails due to
-                    # http_put having a side effect of "put_error"
-                    #
-                    # The agent tries to upload using a direct connection, and that succeeds.
-                    self.assertEqual(1, wire_protocol_client.status_blob.upload.call_count)
-                    # The agent never touches the default protocol is this code path, so no change.
-                    self.assertFalse(wire.HostPluginProtocol.is_default_channel())
-                    # The agent never logs telemetry event for direct fallback
-                    self.assertEqual(1, patch_add_event.call_count)
-                    self.assertEqual('ReportStatus', patch_add_event.call_args[1]['op'])
-                    self.assertTrue('Falling back to direct' in patch_add_event.call_args[1]['message'])
-                    self.assertEqual(True, patch_add_event.call_args[1]['is_success'])
+            with self.create_mock_protocol() as wire_protocol:
+                wire_protocol_client = wire_protocol.client
+
+                put_error = wire.HttpError("put status http error")
+                with patch.object(restutil, "http_put", side_effect=put_error):
+                    with patch.object(wire.HostPluginProtocol,
+                                      "ensure_initialized", return_value=True):
+                        self.assertRaises(wire.ProtocolError, wire_protocol_client.upload_status_blob)
+
+                        # The agent tries to upload via HostPlugin and that fails due to
+                        # http_put having a side effect of "put_error"
+                        #
+                        # The agent tries to upload using a direct connection, and that succeeds.
+                        self.assertEqual(1, wire_protocol_client.status_blob.upload.call_count)
+                        # The agent never touches the default protocol is this code path, so no change.
+                        self.assertFalse(wire.HostPluginProtocol.is_default_channel())
+                        # The agent never logs telemetry event for direct fallback
+                        self.assertEqual(1, patch_add_event.call_count)
+                        self.assertEqual('ReportStatus', patch_add_event.call_args[1]['op'])
+                        self.assertTrue('Falling back to direct' in patch_add_event.call_args[1]['message'])
+                        self.assertEqual(True, patch_add_event.call_args[1]['is_success'])
 
     def test_validate_http_request(self):
         """Validate correct set of data is sent to HostGAPlugin when reporting VM status"""
 
-        wire_protocol_client = wire.WireProtocol(wireserver_url).client
-        test_goal_state = wire.GoalState(WireProtocolData(DATA_FILE).goal_state)
-
-        status_blob = wire_protocol_client.status_blob
-        status_blob.data = faux_status
-        status_blob.vm_status = restapi.VMStatus(message="Ready", status="Ready")
+        with mock_wire_protocol.create(DATA_FILE) as protocol:
+            test_goal_state = protocol.client._goal_state
 
-        exp_method = 'PUT'
-        exp_url = hostplugin_status_url
-        exp_data = self._hostplugin_data(
-            status_blob.get_block_blob_headers(len(faux_status)),
-            bytearray(faux_status, encoding='utf-8'))
+            status_blob = protocol.client.status_blob
+            status_blob.data = faux_status
+            status_blob.vm_status = restapi.VMStatus(message="Ready", status="Ready")
 
-        with patch.object(restutil, "http_request") as patch_http:
-            patch_http.return_value = Mock(status=httpclient.OK)
+            exp_method = 'PUT'
+            exp_url = hostplugin_status_url
+            exp_data = self._hostplugin_data(
+                status_blob.get_block_blob_headers(len(faux_status)),
+                bytearray(faux_status, encoding='utf-8'))
 
-            wire_protocol_client.get_goal_state = Mock(return_value=test_goal_state)
-            plugin = wire_protocol_client.get_host_plugin()
+            with patch.object(restutil, "http_request") as patch_http:
+                patch_http.return_value = Mock(status=httpclient.OK)
 
-            with patch.object(plugin, 'get_api_versions') as patch_api:
-                patch_api.return_value = API_VERSION
-                plugin.put_vm_status(status_blob, sas_url, block_blob_type)
+                protocol.client.get_goal_state = Mock(return_value=test_goal_state)
+                plugin = protocol.client.get_host_plugin()
 
-                self.assertTrue(patch_http.call_count == 2)
+                with patch.object(plugin, 'get_api_versions') as patch_api:
+                    patch_api.return_value = API_VERSION
+                    plugin.put_vm_status(status_blob, sas_url, block_blob_type)
 
-                # first call is to host plugin
-                self._validate_hostplugin_args(
-                    patch_http.call_args_list[0],
-                    test_goal_state,
-                    exp_method, exp_url, exp_data)
+                    self.assertTrue(patch_http.call_count == 2)
 
-                # second call is to health service
-                self.assertEqual('POST', patch_http.call_args_list[1][0][0])
-                self.assertEqual(health_service_url, patch_http.call_args_list[1][0][1])
+                    # first call is to host plugin
+                    self._validate_hostplugin_args(
+                        patch_http.call_args_list[0],
+                        test_goal_state,
+                        exp_method, exp_url, exp_data)
 
-    @patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state")
-    def test_no_fallback(self, _):
-        """
-        Validate fallback to upload status using HostGAPlugin is not happening
-        when status reporting via default method is successful
-        """
-        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
-                ext_config = wire.ExtensionsConfig(None)
-                ext_config.status_upload_blob = sas_url
-                wire_protocol_client._set_ext_conf(ext_config)
-                wire_protocol_client.status_blob.vm_status = vmstatus
-                wire_protocol_client.upload_status_blob()
-
-                self.assertTrue(patch_put.call_count == 0, "Fallback was engaged")
+                    # second call is to health service
+                    self.assertEqual('POST', patch_http.call_args_list[1][0][0])
+                    self.assertEqual(health_service_url, patch_http.call_args_list[1][0][1])
 
     def test_validate_block_blob(self):
-        """Validate correct set of data is sent to HostGAPlugin when reporting VM status"""
-        wire_protocol_client = wire.WireProtocol(wireserver_url).client
-        test_goal_state = wire.GoalState(WireProtocolData(DATA_FILE).goal_state)
-
-        host_client = wire.HostPluginProtocol(wireserver_url,
-                                              test_goal_state.container_id,
-                                              test_goal_state.role_config_name)
-        self.assertFalse(host_client.is_initialized)
-        self.assertTrue(host_client.api_versions is None)
-        self.assertTrue(host_client.health_service is not None)
-
-        status_blob = wire_protocol_client.status_blob
-        status_blob.data = faux_status
-        status_blob.type = block_blob_type
-        status_blob.vm_status = restapi.VMStatus(message="Ready", status="Ready")
-
-        exp_method = 'PUT'
-        exp_url = hostplugin_status_url
-        exp_data = self._hostplugin_data(
-            status_blob.get_block_blob_headers(len(faux_status)),
-            bytearray(faux_status, encoding='utf-8'))
-
-        with patch.object(restutil, "http_request") as patch_http:
-            patch_http.return_value = Mock(status=httpclient.OK)
+        with mock_wire_protocol.create(DATA_FILE) as protocol:
+            test_goal_state = protocol.client._goal_state
+
+            host_client = wire.HostPluginProtocol(wireserver_url,
+                                                  test_goal_state.container_id,
+                                                  test_goal_state.role_config_name)
+            self.assertFalse(host_client.is_initialized)
+            self.assertTrue(host_client.api_versions is None)
+            self.assertTrue(host_client.health_service is not None)
+
+            status_blob = protocol.client.status_blob
+            status_blob.data = faux_status
+            status_blob.type = block_blob_type
+            status_blob.vm_status = restapi.VMStatus(message="Ready", status="Ready")
+
+            exp_method = 'PUT'
+            exp_url = hostplugin_status_url
+            exp_data = self._hostplugin_data(
+                status_blob.get_block_blob_headers(len(faux_status)),
+                bytearray(faux_status, encoding='utf-8'))
+
+            with patch.object(restutil, "http_request") as patch_http:
+                patch_http.return_value = Mock(status=httpclient.OK)
 
-            with patch.object(wire.HostPluginProtocol,
-                              "get_api_versions") as patch_get:
-                patch_get.return_value = api_versions
-                host_client.put_vm_status(status_blob, sas_url)
+                with patch.object(wire.HostPluginProtocol,
+                                  "get_api_versions") as patch_get:
+                    patch_get.return_value = api_versions
+                    host_client.put_vm_status(status_blob, sas_url)
 
-                self.assertTrue(patch_http.call_count == 2)
+                    self.assertTrue(patch_http.call_count == 2)
 
-                # first call is to host plugin
-                self._validate_hostplugin_args(
-                    patch_http.call_args_list[0],
-                    test_goal_state,
-                    exp_method, exp_url, exp_data)
+                    # first call is to host plugin
+                    self._validate_hostplugin_args(
+                        patch_http.call_args_list[0],
+                        test_goal_state,
+                        exp_method, exp_url, exp_data)
 
-                # second call is to health service
-                self.assertEqual('POST', patch_http.call_args_list[1][0][0])
-                self.assertEqual(health_service_url, patch_http.call_args_list[1][0][1])
+                    # second call is to health service
+                    self.assertEqual('POST', patch_http.call_args_list[1][0][0])
+                    self.assertEqual(health_service_url, patch_http.call_args_list[1][0][1])
 
     def test_validate_page_blobs(self):
         """Validate correct set of data is sent for page blobs"""
-        wire_protocol_client = wire.WireProtocol(wireserver_url).client
-        test_goal_state = wire.GoalState(WireProtocolData(DATA_FILE).goal_state)
+        with mock_wire_protocol.create(DATA_FILE) as protocol:
+            test_goal_state = protocol.client._goal_state
 
-        host_client = wire.HostPluginProtocol(wireserver_url,
-                                              test_goal_state.container_id,
-                                              test_goal_state.role_config_name)
+            host_client = wire.HostPluginProtocol(wireserver_url,
+                                                  test_goal_state.container_id,
+                                                  test_goal_state.role_config_name)
 
-        self.assertFalse(host_client.is_initialized)
-        self.assertTrue(host_client.api_versions is None)
+            self.assertFalse(host_client.is_initialized)
+            self.assertTrue(host_client.api_versions is None)
 
-        status_blob = wire_protocol_client.status_blob
-        status_blob.data = faux_status
-        status_blob.type = page_blob_type
-        status_blob.vm_status = restapi.VMStatus(message="Ready", status="Ready")
+            status_blob = protocol.client.status_blob
+            status_blob.data = faux_status
+            status_blob.type = page_blob_type
+            status_blob.vm_status = restapi.VMStatus(message="Ready", status="Ready")
+
+            exp_method = 'PUT'
+            exp_url = hostplugin_status_url
 
-        exp_method = 'PUT'
-        exp_url = hostplugin_status_url
-
-        page_status = bytearray(status_blob.data, encoding='utf-8')
-        page_size = int((len(page_status) + 511) / 512) * 512
-        page_status = bytearray(status_blob.data.ljust(page_size), encoding='utf-8')
-        page = bytearray(page_size)
-        page[0: page_size] = page_status[0: len(page_status)]
-        mock_response = MockResponse('', httpclient.OK)
-
-        with patch.object(restutil, "http_request",
-                          return_value=mock_response) as patch_http:
-            with patch.object(wire.HostPluginProtocol,
-                              "get_api_versions") as patch_get:
-                patch_get.return_value = api_versions
-                host_client.put_vm_status(status_blob, sas_url)
-
-                self.assertTrue(patch_http.call_count == 3)
-
-                # first call is to host plugin
-                exp_data = self._hostplugin_data(
-                    status_blob.get_page_blob_create_headers(
-                        page_size))
-                self._validate_hostplugin_args(
-                    patch_http.call_args_list[0],
-                    test_goal_state,
-                    exp_method, exp_url, exp_data)
-
-                # second call is to health service
-                self.assertEqual('POST', patch_http.call_args_list[1][0][0])
-                self.assertEqual(health_service_url, patch_http.call_args_list[1][0][1])
-
-                # last call is to host plugin
-                exp_data = self._hostplugin_data(
-                    status_blob.get_page_blob_page_headers(
-                        0, page_size),
-                    page)
-                exp_data['requestUri'] += "?comp=page"
-                self._validate_hostplugin_args(
-                    patch_http.call_args_list[2],
-                    test_goal_state,
-                    exp_method, exp_url, exp_data)
+            page_status = bytearray(status_blob.data, encoding='utf-8')
+            page_size = int((len(page_status) + 511) / 512) * 512
+            page_status = bytearray(status_blob.data.ljust(page_size), encoding='utf-8')
+            page = bytearray(page_size)
+            page[0: page_size] = page_status[0: len(page_status)]
+            mock_response = MockResponse('', httpclient.OK)
+
+            with patch.object(restutil, "http_request",
+                              return_value=mock_response) as patch_http:
+                with patch.object(wire.HostPluginProtocol,
+                                  "get_api_versions") as patch_get:
+                    patch_get.return_value = api_versions
+                    host_client.put_vm_status(status_blob, sas_url)
+
+                    self.assertTrue(patch_http.call_count == 3)
+
+                    # first call is to host plugin
+                    exp_data = self._hostplugin_data(
+                        status_blob.get_page_blob_create_headers(
+                            page_size))
+                    self._validate_hostplugin_args(
+                        patch_http.call_args_list[0],
+                        test_goal_state,
+                        exp_method, exp_url, exp_data)
+
+                    # second call is to health service
+                    self.assertEqual('POST', patch_http.call_args_list[1][0][0])
+                    self.assertEqual(health_service_url, patch_http.call_args_list[1][0][1])
+
+                    # last call is to host plugin
+                    exp_data = self._hostplugin_data(
+                        status_blob.get_page_blob_page_headers(
+                            0, page_size),
+                        page)
+                    exp_data['requestUri'] += "?comp=page"
+                    self._validate_hostplugin_args(
+                        patch_http.call_args_list[2],
+                        test_goal_state,
+                        exp_method, exp_url, exp_data)
 
     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_config_name,
-                            "x-ms-artifact-location": sas_url}
-
-        host_client = wire.HostPluginProtocol(wireserver_url,
-                                              test_goal_state.container_id,
-                                              test_goal_state.role_config_name)
-        self.assertFalse(host_client.is_initialized)
-        self.assertTrue(host_client.api_versions is None)
-        self.assertTrue(host_client.health_service is not None)
-
-        with patch.object(wire.HostPluginProtocol, "get_api_versions", return_value=api_versions) as patch_get:
-            actual_url, actual_headers = host_client.get_artifact_request(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])
+        with mock_wire_protocol.create(DATA_FILE) as protocol:
+            test_goal_state = protocol.client._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_config_name,
+                                "x-ms-artifact-location": sas_url}
+
+            host_client = wire.HostPluginProtocol(wireserver_url,
+                                                  test_goal_state.container_id,
+                                                  test_goal_state.role_config_name)
+            self.assertFalse(host_client.is_initialized)
+            self.assertTrue(host_client.api_versions is None)
+            self.assertTrue(host_client.health_service is not None)
+
+            with patch.object(wire.HostPluginProtocol, "get_api_versions", return_value=api_versions) as patch_get:
+                actual_url, actual_headers = host_client.get_artifact_request(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])
 
     @patch("azurelinuxagent.common.utils.restutil.http_get")
     def test_health(self, patch_http_get):
diff --git a/tests/protocol/test_protocol_util.py b/tests/protocol/test_protocol_util.py
index c3dc6d04..1107c280 100644
--- a/tests/protocol/test_protocol_util.py
+++ b/tests/protocol/test_protocol_util.py
@@ -15,14 +15,13 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
-import os
 import unittest
 from multiprocessing import Queue
 from threading import Thread
 
 from tests.tools import AgentTestCase, MagicMock, Mock, patch, clear_singleton_instances
 from azurelinuxagent.common.exception import *
-from azurelinuxagent.common.protocol import get_protocol_util
+from azurelinuxagent.common.protocol.util import get_protocol_util
 from azurelinuxagent.common.utils.restutil import KNOWN_WIRESERVER_IP
 from azurelinuxagent.common.protocol.util import ProtocolUtil
 from errno import ENOENT
diff --git a/tests/protocol/test_wire.py b/tests/protocol/test_wire.py
index 995c1a00..84cbfc95 100644
--- a/tests/protocol/test_wire.py
+++ b/tests/protocol/test_wire.py
@@ -18,26 +18,25 @@
 
 import json
 import os
-import stat
-import tempfile
 import time
 import unittest
 import uuid
-import zipfile
+import contextlib
 
 from azurelinuxagent.common.exception import InvalidContainerError, ResourceGoneError, ProtocolError, \
     ExtensionDownloadError
 from azurelinuxagent.common.future import httpclient
 from azurelinuxagent.common.protocol.hostplugin import HostPluginProtocol
-from azurelinuxagent.common.protocol.wire import WireProtocol, WireClient, GoalState, ExtensionsConfig, \
-    InVMArtifactsProfile, VMAgentManifestUri, StatusBlob, VMStatus, ExtHandlerVersionUri, DataContractList, socket
+from azurelinuxagent.common.protocol.goal_state import GoalState, ExtensionsConfig
+from azurelinuxagent.common.protocol.wire import WireProtocol, WireClient, InVMArtifactsProfile, VMAgentManifestUri, \
+    StatusBlob, VMStatus, ExtHandlerVersionUri, DataContractList, socket
 from azurelinuxagent.common.telemetryevent import TelemetryEvent, TelemetryEventParam, TelemetryEventList
-from azurelinuxagent.common.utils import fileutil, restutil
-from azurelinuxagent.common.utils.shellutil import run_get_output
+from azurelinuxagent.common.utils import restutil
 from azurelinuxagent.common.version import CURRENT_VERSION, DISTRO_NAME, DISTRO_VERSION
 from tests.ga.test_monitor import random_generator
-from tests.protocol import mockwiredata
-from tests.protocol.mockwiredata import WireProtocolData
+from tests.protocol import mockwiredata, mock_wire_protocol
+from tests.protocol.mockwiredata import WireProtocolData, DATA_FILE_NO_EXT
+
 from tests.tools import ANY, MagicMock, Mock, patch, AgentTestCase, skip_if_predicate_true, running_under_travis
 
 data_with_bom = b'\xef\xbb\xbfhehe'
@@ -60,14 +59,18 @@ def get_event(message, duration=30000, evt_type="", is_internal=False, is_succes
     return event
 
 
-def set_up_wireclient_ext_conf(wire_protocol_client, artifacts_profile_blob=None,
-                               status_upload_blob=None,
-                               status_upload_blob_type=None):
-    ext_conf = ExtensionsConfig(None)
-    ext_conf.artifacts_profile_blob = artifacts_profile_blob
-    ext_conf.status_upload_blob = status_upload_blob
-    ext_conf.status_upload_blob_type = status_upload_blob_type
-    wire_protocol_client._set_ext_conf(ext_conf)
+@contextlib.contextmanager
+def create_mock_protocol(artifacts_profile_blob=None, status_upload_blob=None, status_upload_blob_type=None):
+    with mock_wire_protocol.create(DATA_FILE_NO_EXT) as protocol:
+        # These tests use mock wire data that dont have any extensions (extension config will be empty).
+        # Populate the upload blob and artifacts profile blob.
+        ext_conf = ExtensionsConfig(None)
+        ext_conf.artifacts_profile_blob = artifacts_profile_blob
+        ext_conf.status_upload_blob = status_upload_blob
+        ext_conf.status_upload_blob_type = status_upload_blob_type
+        protocol.client._goal_state.ext_conf = ext_conf
+
+        yield protocol
 
 
 @patch("time.sleep")
@@ -192,175 +195,129 @@ class TestWireProtocol(AgentTestCase):
                 self.assertTrue(c == (True if i != 3 else False))
 
     def test_status_blob_parsing(self, *args):
-        wire_protocol_client = WireProtocol(WIRESERVER_URL).client
-        wire_protocol_client._set_ext_conf(
-            ExtensionsConfig(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).ext_conf))
-        self.assertEqual(wire_protocol_client.get_ext_conf().status_upload_blob,
-                         'https://test.blob.core.windows.net/vhds/test-cs12.test-cs12.test-cs12.status?'
-                         'sr=b&sp=rw&se=9999-01-01&sk=key1&sv=2014-02-14&'
-                         'sig=hfRh7gzUE7sUtYwke78IOlZOrTRCYvkec4hGZ9zZzXo')
-        self.assertEqual(wire_protocol_client.get_ext_conf().status_upload_blob_type,
-                         u'BlockBlob')
-        pass
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE) as protocol:
+            self.assertEqual(protocol.client.get_ext_conf().status_upload_blob,
+                             'https://test.blob.core.windows.net/vhds/test-cs12.test-cs12.test-cs12.status?'
+                             'sr=b&sp=rw&se=9999-01-01&sk=key1&sv=2014-02-14&'
+                             'sig=hfRh7gzUE7sUtYwke78IOlZOrTRCYvkec4hGZ9zZzXo')
+            self.assertEqual(protocol.client.get_ext_conf().status_upload_blob_type, u'BlockBlob')
 
     def test_get_host_ga_plugin(self, *args):
-        wire_protocol_client = WireProtocol(WIRESERVER_URL).client
-        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.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()
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE) as protocol:
+            host_plugin = protocol.client.get_host_plugin()
+            goal_state = protocol.client.get_goal_state()
             self.assertEqual(goal_state.container_id, host_plugin.container_id)
             self.assertEqual(goal_state.role_config_name, host_plugin.role_config_name)
-            self.assertEqual(1, patch_get_goal_state.call_count)
 
-    @patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state")
     def test_upload_status_blob_default(self, *args):
         """
         Default status blob method is HostPlugin.
         """
-        vmstatus = VMStatus(message="Ready", status="Ready")
-        wire_protocol_client = WireProtocol(WIRESERVER_URL).client
-        set_up_wireclient_ext_conf(wire_protocol_client, status_upload_blob=testurl, status_upload_blob_type=testtype)
-        wire_protocol_client.status_blob.vm_status = vmstatus
+        with create_mock_protocol(status_upload_blob=testurl, status_upload_blob_type=testtype) as protocol:
+            protocol.client.status_blob.vm_status = VMStatus(message="Ready", status="Ready")
+
+            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") as patch_default_upload:
+                        HostPluginProtocol.set_default_channel(False)
+                        protocol.client.upload_status_blob()
+
+                        # do not call the direct method unless host plugin fails
+                        patch_default_upload.assert_not_called()
+                        # host plugin always fetches a goal state
+                        patch_get_goal_state.assert_called_once_with()
+                        # host plugin uploads the status blob
+                        patch_host_ga_plugin_upload.assert_called_once_with(ANY, testurl, 'BlockBlob')
+
+    def test_upload_status_blob_host_ga_plugin(self, *_):
+        with create_mock_protocol(status_upload_blob=testurl, status_upload_blob_type=testtype) as protocol:
+            protocol.client.status_blob.vm_status = VMStatus(message="Ready", status="Ready")
+
+            with patch.object(HostPluginProtocol, "ensure_initialized", return_value=True):
+                with patch.object(StatusBlob, "upload", return_value=False) as patch_default_upload:
+                    with patch.object(HostPluginProtocol, "_put_block_blob_status") as patch_http:
+                        HostPluginProtocol.set_default_channel(False)
+                        protocol.client.upload_status_blob()
+                        patch_default_upload.assert_not_called()
+                        patch_http.assert_called_once_with(testurl, protocol.client.status_blob)
+                        self.assertFalse(HostPluginProtocol.is_default_channel())
 
-        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") as patch_default_upload:
-                    HostPluginProtocol.set_default_channel(False)
-                    wire_protocol_client.upload_status_blob()
-
-                    # do not call the direct method unless host plugin fails
-                    patch_default_upload.assert_not_called()
-                    # host plugin always fetches a goal state
-                    patch_get_goal_state.assert_called_once_with()
-                    # host plugin uploads the status blob
-                    patch_host_ga_plugin_upload.assert_called_once_with(ANY, testurl, 'BlockBlob')
-
-    @patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state")
-    def test_upload_status_blob_host_ga_plugin(self, *args):
-        vmstatus = VMStatus(message="Ready", status="Ready")
-        wire_protocol_client = WireProtocol(WIRESERVER_URL).client
-        set_up_wireclient_ext_conf(wire_protocol_client, status_upload_blob=testurl, status_upload_blob_type=testtype)
-        wire_protocol_client.status_blob.vm_status = vmstatus
-        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
-
-        with patch.object(HostPluginProtocol,
-                          "ensure_initialized",
-                          return_value=True):
-            with patch.object(StatusBlob,
-                              "upload",
-                              return_value=False) as patch_default_upload:
-                with patch.object(HostPluginProtocol,
-                                  "_put_block_blob_status") as patch_http:
-                    HostPluginProtocol.set_default_channel(False)
-                    wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
-                    wire_protocol_client.upload_status_blob()
-                    patch_default_upload.assert_not_called()
-                    self.assertEqual(1, wire_protocol_client.get_goal_state.call_count)
-                    patch_http.assert_called_once_with(testurl, wire_protocol_client.status_blob)
-                    self.assertFalse(HostPluginProtocol.is_default_channel())
-
-    @patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state")
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.ensure_initialized")
-    def test_upload_status_blob_unknown_type_assumes_block(self, _, __, *args):
-        vmstatus = VMStatus(message="Ready", status="Ready")
-        wire_protocol_client = WireProtocol(WIRESERVER_URL).client
-        set_up_wireclient_ext_conf(wire_protocol_client,
-                                   status_upload_blob=testurl,
-                                   status_upload_blob_type="NotALegalType")
-        wire_protocol_client.status_blob.vm_status = vmstatus
-
-        with patch.object(WireClient, "get_goal_state") as patch_get_goal_state:
+    def test_upload_status_blob_unknown_type_assumes_block(self, *_):
+        with create_mock_protocol(status_upload_blob=testurl, status_upload_blob_type="NotALegalType") as protocol:
+            protocol.client.status_blob.vm_status = VMStatus(message="Ready", status="Ready")
+
             with patch.object(StatusBlob, "prepare") as patch_prepare:
                 with patch.object(StatusBlob, "upload") as patch_default_upload:
                     HostPluginProtocol.set_default_channel(False)
-                    wire_protocol_client.upload_status_blob()
+                    protocol.client.upload_status_blob()
 
                     patch_prepare.assert_called_once_with("BlockBlob")
                     patch_default_upload.assert_called_once_with(testurl)
-                    patch_get_goal_state.assert_called_once_with()
-
-    @patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state")
-    def test_upload_status_blob_reports_prepare_error(self, *args):
-        vmstatus = VMStatus(message="Ready", status="Ready")
-        wire_protocol_client = WireProtocol(WIRESERVER_URL).client
-        set_up_wireclient_ext_conf(wire_protocol_client, status_upload_blob=testurl, status_upload_blob_type=testtype)
-        wire_protocol_client.status_blob.vm_status = vmstatus
-        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
 
-        with patch.object(StatusBlob, "prepare", side_effect=Exception) as mock_prepare:
-            self.assertRaises(ProtocolError, wire_protocol_client.upload_status_blob)
-            self.assertEqual(1, mock_prepare.call_count)
+    def test_upload_status_blob_reports_prepare_error(self, *_):
+        with create_mock_protocol(status_upload_blob=testurl, status_upload_blob_type=testtype) as protocol:
+            protocol.client.status_blob.vm_status = VMStatus(message="Ready", status="Ready")
 
-    def test_get_in_vm_artifacts_profile_blob_not_available(self, *args):
-        wire_protocol_client = WireProtocol(WIRESERVER_URL).client
-        wire_protocol_client._set_ext_conf(ExtensionsConfig(None))
+            with patch.object(StatusBlob, "prepare", side_effect=Exception) as mock_prepare:
+                self.assertRaises(ProtocolError, protocol.client.upload_status_blob)
+                self.assertEqual(1, mock_prepare.call_count)
 
+    def test_get_in_vm_artifacts_profile_blob_not_available(self, *_):
         # Test when artifacts_profile_blob is null/None
-        self.assertEqual(None, wire_protocol_client.get_artifacts_profile())
+        with mock_wire_protocol.create(DATA_FILE_NO_EXT) as protocol:
+            protocol.client._goal_state.ext_conf = ExtensionsConfig(None)
+
+            self.assertEqual(None, protocol.client.get_artifacts_profile())
 
         # Test when artifacts_profile_blob is whitespace
-        set_up_wireclient_ext_conf(wire_protocol_client, artifacts_profile_blob="  ")
-        self.assertEqual(None, wire_protocol_client.get_artifacts_profile())
-
-    def test_get_in_vm_artifacts_profile_response_body_not_valid(self, *args):
-        wire_protocol_client = WireProtocol(WIRESERVER_URL).client
-        set_up_wireclient_ext_conf(wire_protocol_client, artifacts_profile_blob=testurl)
-        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
-        wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
-
-        with patch.object(HostPluginProtocol, "get_artifact_request",
-                          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_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_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_artifacts_profile()
-            self.assertEqual(dict(), in_vm_artifacts_profile.__dict__,
-                             'If 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)
+        with create_mock_protocol(artifacts_profile_blob="  ") as protocol:
+            self.assertEqual(None, protocol.client.get_artifacts_profile())
+
+    def test_get_in_vm_artifacts_profile_response_body_not_valid(self, *_):
+        with create_mock_protocol(artifacts_profile_blob=testurl) as protocol:
+            with patch.object(HostPluginProtocol, "get_artifact_request", return_value=['dummy_url', {}]) as host_plugin_get_artifact_url_and_headers:
+                # Test when response body is None
+                protocol.client.call_storage_service = Mock(return_value=MockResponse(None, 200))
+                in_vm_artifacts_profile = protocol.client.get_artifacts_profile()
+                self.assertTrue(in_vm_artifacts_profile is None)
+
+                # Test when response body is None
+                protocol.client.call_storage_service = Mock(return_value=MockResponse('   '.encode('utf-8'), 200))
+                in_vm_artifacts_profile = protocol.client.get_artifacts_profile()
+                self.assertTrue(in_vm_artifacts_profile is None)
+
+                # Test when response body is None
+                protocol.client.call_storage_service = Mock(return_value=MockResponse('{ }'.encode('utf-8'), 200))
+                in_vm_artifacts_profile = protocol.client.get_artifacts_profile()
+                self.assertEqual(dict(), in_vm_artifacts_profile.__dict__,
+                                 'If 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)
 
     @patch("azurelinuxagent.common.event.add_event")
     def test_artifacts_profile_json_parsing(self, patch_event, *args):
-        wire_protocol_client = WireProtocol(WIRESERVER_URL).client
-        set_up_wireclient_ext_conf(wire_protocol_client, artifacts_profile_blob=testurl)
-        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
-        wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
-
-        # response is invalid json
-        wire_protocol_client.call_storage_service = Mock(return_value=MockResponse("invalid json".encode('utf-8'), 200))
-        in_vm_artifacts_profile = wire_protocol_client.get_artifacts_profile()
+        with create_mock_protocol(artifacts_profile_blob=testurl) as protocol:
+            # response is invalid json
+            protocol.client.call_storage_service = Mock(return_value=MockResponse("invalid json".encode('utf-8'), 200))
+            in_vm_artifacts_profile = protocol.client.get_artifacts_profile()
 
-        # ensure response is empty
-        self.assertEqual(None, in_vm_artifacts_profile)
+            # ensure response is empty
+            self.assertEqual(None, in_vm_artifacts_profile)
 
-        # ensure event is logged
-        self.assertEqual(1, patch_event.call_count)
-        self.assertFalse(patch_event.call_args[1]['is_success'])
-        self.assertTrue('invalid json' in patch_event.call_args[1]['message'])
-        self.assertEqual('ArtifactsProfileBlob', patch_event.call_args[1]['op'])
+            # ensure event is logged
+            self.assertEqual(1, patch_event.call_count)
+            self.assertFalse(patch_event.call_args[1]['is_success'])
+            self.assertTrue('invalid json' in patch_event.call_args[1]['message'])
+            self.assertEqual('ArtifactsProfileBlob', patch_event.call_args[1]['op'])
 
     def test_get_in_vm_artifacts_profile_default(self, *args):
-        wire_protocol_client = WireProtocol(WIRESERVER_URL).client
-        set_up_wireclient_ext_conf(wire_protocol_client, artifacts_profile_blob=testurl)
-        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.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_artifacts_profile()
-        self.assertEqual(dict(onHold='true'), in_vm_artifacts_profile.__dict__)
-        self.assertTrue(in_vm_artifacts_profile.is_on_hold())
+        with create_mock_protocol(artifacts_profile_blob=testurl) as protocol:
+            protocol.client.call_storage_service = Mock(return_value=MockResponse('{"onHold": "true"}'.encode('utf-8'), 200))
+            in_vm_artifacts_profile = protocol.client.get_artifacts_profile()
+            self.assertEqual(dict(onHold='true'), in_vm_artifacts_profile.__dict__)
+            self.assertTrue(in_vm_artifacts_profile.is_on_hold())
 
     def test_fetch_manifest_fallback(self, *args):
         uri1 = ExtHandlerVersionUri()
@@ -415,19 +372,14 @@ class TestWireProtocol(AgentTestCase):
                 self.assertTrue(HostPluginProtocol.is_default_channel())
 
     def test_get_in_vm_artifacts_profile_host_ga_plugin(self, *args):
-        wire_protocol_client = WireProtocol(WIRESERVER_URL).client
-        set_up_wireclient_ext_conf(wire_protocol_client, artifacts_profile_blob=testurl)
-        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
-        wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
-        wire_protocol_client.fetch = Mock(side_effect=[None, '{"onHold": "true"}'])
-        with patch.object(HostPluginProtocol,
-                          "get_artifact_request",
-                          return_value=['dummy_url', {}]) as artifact_request:
-            in_vm_artifacts_profile = wire_protocol_client.get_artifacts_profile()
-            self.assertTrue(in_vm_artifacts_profile is not None)
-            self.assertEqual(dict(onHold='true'), in_vm_artifacts_profile.__dict__)
-            self.assertTrue(in_vm_artifacts_profile.is_on_hold())
-            artifact_request.assert_called_once_with(testurl)
+        with create_mock_protocol(artifacts_profile_blob=testurl) as protocol:
+            protocol.client.fetch = Mock(side_effect=[None, '{"onHold": "true"}'])
+            with patch.object(HostPluginProtocol, "get_artifact_request", return_value=['dummy_url', {}]) as artifact_request:
+                in_vm_artifacts_profile = protocol.client.get_artifacts_profile()
+                self.assertTrue(in_vm_artifacts_profile is not None)
+                self.assertEqual(dict(onHold='true'), in_vm_artifacts_profile.__dict__)
+                self.assertTrue(in_vm_artifacts_profile.is_on_hold())
+                artifact_request.assert_called_once_with(testurl)
 
     @patch("socket.gethostname", return_value="hostname")
     @patch("time.gmtime", return_value=time.localtime(1485543256))
@@ -535,20 +487,14 @@ class TestWireProtocol(AgentTestCase):
 class TestWireClient(AgentTestCase):
 
     def test_get_ext_conf_without_uri(self, *args):
-        test_goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
-        test_goal_state.ext_uri = None
-        wire_protocol_client = WireProtocol(WIRESERVER_URL).client
-        wire_protocol_client.get_goal_state = Mock(return_value=test_goal_state)
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE_NO_EXT) as protocol:
+            ext_conf = protocol.client.get_ext_conf()
 
-        with patch("azurelinuxagent.common.protocol.wire.WireClient.fetch_cache",
-                   return_value=mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).ext_conf):
-            ext_conf = wire_protocol_client.get_ext_conf()
-
-        self.assertEqual(0, len(ext_conf.ext_handlers.extHandlers))
-        self.assertEqual(0, len(ext_conf.vmagent_manifests.vmAgentManifests))
-        self.assertIsNone(ext_conf.status_upload_blob)
-        self.assertIsNone(ext_conf.status_upload_blob_type)
-        self.assertIsNone(ext_conf.artifacts_profile_blob)
+            self.assertEqual(0, len(ext_conf.ext_handlers.extHandlers))
+            self.assertEqual(0, len(ext_conf.vmagent_manifests.vmAgentManifests))
+            self.assertIsNone(ext_conf.status_upload_blob)
+            self.assertIsNone(ext_conf.status_upload_blob_type)
+            self.assertIsNone(ext_conf.artifacts_profile_blob)
 
     def test_get_ext_conf_with_uri(self, *args):
         mock_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
@@ -811,65 +757,56 @@ class TestWireClient(AgentTestCase):
 
                             self.assertEquals(HostPluginProtocol.is_default_channel(), False)
 
-    @patch("azurelinuxagent.common.protocol.wire.WireClient.get_goal_state")
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.get_artifact_request")
-    def test_get_artifacts_profile_should_not_invoke_host_channel_when_direct_channel_succeeds(self,
-                                                                                               mock_get_artifact_request,
-                                                                                               *args):
+    def test_get_artifacts_profile_should_not_invoke_host_channel_when_direct_channel_succeeds(self, mock_get_artifact_request):
         mock_get_artifact_request.return_value = "dummy_url", "dummy_header"
-        client = WireClient("foo.bar")
-        set_up_wireclient_ext_conf(client, artifacts_profile_blob="testurl")
-        json_profile = b'{ "onHold": true }'
+        mock_successful_response = MockResponse(body=b'{ "onHold": true }', status_code=200)
 
         HostPluginProtocol.set_default_channel(False)
-        mock_successful_response = MockResponse(body=json_profile, status_code=200)
 
         # Direct channel succeeds
-        with patch("azurelinuxagent.common.utils.restutil._http_request", return_value=mock_successful_response):
-            with patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state") as mock_update_goal_state:
-                with patch("azurelinuxagent.common.protocol.wire.WireClient.fetch", wraps=client.fetch) as patch_direct:
-                    with patch("azurelinuxagent.common.protocol.wire.WireClient.get_artifacts_profile_through_host",
-                               wraps=client.get_artifacts_profile_through_host) as patch_host:
-                        ret = client.get_artifacts_profile()
-                        self.assertIsInstance(ret, InVMArtifactsProfile)
-
-                        self.assertEquals(patch_host.call_count, 0)
-                        self.assertEquals(patch_direct.call_count, 1)
-                        self.assertEquals(mock_update_goal_state.call_count, 0)
+        with create_mock_protocol(artifacts_profile_blob="testurl") as protocol:
+            with patch("azurelinuxagent.common.utils.restutil._http_request", return_value=mock_successful_response):
+                with patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state") as mock_update_goal_state:
+                    with patch("azurelinuxagent.common.protocol.wire.WireClient.fetch", wraps=protocol.client.fetch) as patch_direct:
+                        with patch("azurelinuxagent.common.protocol.wire.WireClient.get_artifacts_profile_through_host",
+                                   wraps=protocol.client.get_artifacts_profile_through_host) as patch_host:
+                            ret = protocol.client.get_artifacts_profile()
+                            self.assertIsInstance(ret, InVMArtifactsProfile)
+
+                            self.assertEquals(patch_host.call_count, 0)
+                            self.assertEquals(patch_direct.call_count, 1)
+                            self.assertEquals(mock_update_goal_state.call_count, 0)
 
-                        self.assertEquals(HostPluginProtocol.is_default_channel(), False)
+                            self.assertEquals(HostPluginProtocol.is_default_channel(), False)
 
     @patch("azurelinuxagent.common.protocol.wire.WireClient.get_goal_state")
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.get_artifact_request")
-    def test_get_artifacts_profile_should_use_host_channel_when_direct_channel_fails(self, mock_get_artifact_request,
-                                                                                     *args):
+    def test_get_artifacts_profile_should_use_host_channel_when_direct_channel_fails(self, mock_get_artifact_request, *args):
         mock_get_artifact_request.return_value = "dummy_url", "dummy_header"
-        client = WireClient("foo.bar")
-        set_up_wireclient_ext_conf(client, artifacts_profile_blob="testurl")
-        json_profile = b'{ "onHold": true }'
+        mock_failed_response = MockResponse(body=b"", status_code=httpclient.GONE)
+        mock_successful_response = MockResponse(body=b'{ "onHold": true }', status_code=200)
 
         HostPluginProtocol.set_default_channel(False)
 
-        mock_failed_response = MockResponse(body=b"", status_code=httpclient.GONE)
-        mock_successful_response = MockResponse(body=json_profile, status_code=200)
-
         # Direct channel fails, host channel succeeds. Goal state should not have been updated and host channel
         # should have been set as default
-        with patch("azurelinuxagent.common.utils.restutil._http_request",
-                   side_effect=[mock_failed_response, mock_successful_response]):
-            with patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state") as mock_update_goal_state:
-                with patch("azurelinuxagent.common.protocol.wire.WireClient.fetch", wraps=client.fetch) as patch_direct:
-                    with patch("azurelinuxagent.common.protocol.wire.WireClient.get_artifacts_profile_through_host",
-                               wraps=client.get_artifacts_profile_through_host) as patch_host:
-                        ret = client.get_artifacts_profile()
-                        self.assertIsInstance(ret, InVMArtifactsProfile)
-
-                        self.assertEquals(patch_host.call_count, 1)
-                        # The host channel calls the direct function under the covers
-                        self.assertEquals(patch_direct.call_count, 1 + patch_host.call_count)
-                        self.assertEquals(mock_update_goal_state.call_count, 0)
+        with create_mock_protocol(artifacts_profile_blob="testurl") as protocol:
+            with patch("azurelinuxagent.common.utils.restutil._http_request",
+                       side_effect=[mock_failed_response, mock_successful_response]):
+                with patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state") as mock_update_goal_state:
+                    with patch("azurelinuxagent.common.protocol.wire.WireClient.fetch", wraps=protocol.client.fetch) as patch_direct:
+                        with patch("azurelinuxagent.common.protocol.wire.WireClient.get_artifacts_profile_through_host",
+                                   wraps=protocol.client.get_artifacts_profile_through_host) as patch_host:
+                            ret = protocol.client.get_artifacts_profile()
+                            self.assertIsInstance(ret, InVMArtifactsProfile)
+
+                            self.assertEquals(patch_host.call_count, 1)
+                            # The host channel calls the direct function under the covers
+                            self.assertEquals(patch_direct.call_count, 1 + patch_host.call_count)
+                            self.assertEquals(mock_update_goal_state.call_count, 0)
 
-                        self.assertEquals(HostPluginProtocol.is_default_channel(), True)
+                            self.assertEquals(HostPluginProtocol.is_default_channel(), True)
 
     @patch("azurelinuxagent.common.protocol.wire.WireClient.get_goal_state")
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.get_artifact_request")
@@ -877,335 +814,294 @@ class TestWireClient(AgentTestCase):
                                                                                                   mock_get_artifact_request,
                                                                                                   *args):
         mock_get_artifact_request.return_value = "dummy_url", "dummy_header"
-        client = WireClient("foo.bar")
-        set_up_wireclient_ext_conf(client, artifacts_profile_blob="testurl")
-        json_profile = b'{ "onHold": true }'
+        mock_failed_response = MockResponse(body=b"", status_code=httpclient.GONE)
+        mock_successful_response = MockResponse(body=b'{ "onHold": true }', status_code=200)
 
         HostPluginProtocol.set_default_channel(False)
 
-        mock_failed_response = MockResponse(body=b"", status_code=httpclient.GONE)
-        mock_successful_response = MockResponse(body=json_profile, status_code=200)
-
         # Direct channel fails, host channel fails due to stale goal state, host channel succeeds after refresh.
         # As a consequence, host plugin should have been updated and host channel should have been set as default.
-        with patch("azurelinuxagent.common.utils.restutil._http_request",
-                   side_effect=[mock_failed_response, mock_failed_response, mock_successful_response]):
-            with patch(
-                    "azurelinuxagent.common.protocol.wire.WireClient.update_host_plugin_from_goal_state") as mock_update_host_plugin_from_goal_state:
-                with patch("azurelinuxagent.common.protocol.wire.WireClient.fetch", wraps=client.fetch) as patch_direct:
-                    with patch("azurelinuxagent.common.protocol.wire.WireClient.get_artifacts_profile_through_host",
-                               wraps=client.get_artifacts_profile_through_host) as patch_host:
-                        ret = client.get_artifacts_profile()
-                        self.assertIsInstance(ret, InVMArtifactsProfile)
+        with create_mock_protocol(artifacts_profile_blob="testurl") as protocol:
+            with patch("azurelinuxagent.common.utils.restutil._http_request",
+                       side_effect=[mock_failed_response, mock_failed_response, mock_successful_response]):
+                with patch("azurelinuxagent.common.protocol.wire.WireClient.update_host_plugin_from_goal_state") as mock_update_host_plugin_from_goal_state:
+                    with patch("azurelinuxagent.common.protocol.wire.WireClient.fetch", wraps=protocol.client.fetch) as patch_direct:
+                        with patch("azurelinuxagent.common.protocol.wire.WireClient.get_artifacts_profile_through_host",
+                                   wraps=protocol.client.get_artifacts_profile_through_host) as patch_host:
+                            ret = protocol.client.get_artifacts_profile()
+                            self.assertIsInstance(ret, InVMArtifactsProfile)
 
-                        self.assertEquals(patch_host.call_count, 2)
-                        # The host channel calls the direct function under the covers
-                        self.assertEquals(patch_direct.call_count, 1 + patch_host.call_count)
-                        self.assertEquals(mock_update_host_plugin_from_goal_state.call_count, 1)
+                            self.assertEquals(patch_host.call_count, 2)
+                            # The host channel calls the direct function under the covers
+                            self.assertEquals(patch_direct.call_count, 1 + patch_host.call_count)
+                            self.assertEquals(mock_update_host_plugin_from_goal_state.call_count, 1)
 
-                        self.assertEquals(HostPluginProtocol.is_default_channel(), True)
+                            self.assertEquals(HostPluginProtocol.is_default_channel(), True)
 
-    @patch("azurelinuxagent.common.protocol.wire.WireClient.get_goal_state")
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.get_artifact_request")
-    def test_get_artifacts_profile_should_refresh_the_host_plugin_and_not_change_default_channel_if_host_plugin_fails(
-            self, mock_get_artifact_request, *args):
+    def test_get_artifacts_profile_should_refresh_the_host_plugin_and_not_change_default_channel_if_host_plugin_fails(self, mock_get_artifact_request):
         mock_get_artifact_request.return_value = "dummy_url", "dummy_header"
-        client = WireClient("foo.bar")
-        set_up_wireclient_ext_conf(client, artifacts_profile_blob="testurl")
+        mock_failed_response = MockResponse(body=b"", status_code=httpclient.GONE)
 
         HostPluginProtocol.set_default_channel(False)
 
-        mock_failed_response = MockResponse(body=b"", status_code=httpclient.GONE)
-
         # Everything fails. Goal state should have been updated and host channel should not have been set as default.
-        with patch("azurelinuxagent.common.utils.restutil._http_request", return_value=mock_failed_response):
-            with patch(
-                    "azurelinuxagent.common.protocol.wire.WireClient.update_host_plugin_from_goal_state") as mock_update_host_plugin_from_goal_state:
-                with patch("azurelinuxagent.common.protocol.wire.WireClient.fetch", wraps=client.fetch) as patch_direct:
-                    with patch("azurelinuxagent.common.protocol.wire.WireClient.get_artifacts_profile_through_host",
-                               wraps=client.get_artifacts_profile_through_host) as patch_host:
-                        ret = client.get_artifacts_profile()
-                        self.assertEquals(ret, None)
+        with create_mock_protocol(artifacts_profile_blob="testurl") as protocol:
+            with patch("azurelinuxagent.common.utils.restutil._http_request", return_value=mock_failed_response):
+                with patch("azurelinuxagent.common.protocol.wire.WireClient.update_host_plugin_from_goal_state") as mock_update_host_plugin_from_goal_state:
+                    with patch("azurelinuxagent.common.protocol.wire.WireClient.fetch", wraps=protocol.client.fetch) as patch_direct:
+                        with patch("azurelinuxagent.common.protocol.wire.WireClient.get_artifacts_profile_through_host", wraps=protocol.client.get_artifacts_profile_through_host) as patch_host:
+                            ret = protocol.client.get_artifacts_profile()
+
+                            self.assertEquals(ret, None)
+                            self.assertEquals(patch_host.call_count, 2)
 
-                        self.assertEquals(patch_host.call_count, 2)
-                        # The host channel calls the direct function under the covers
-                        self.assertEquals(patch_direct.call_count, 1 + patch_host.call_count)
-                        self.assertEquals(mock_update_host_plugin_from_goal_state.call_count, 1)
+                            # The host channel calls the direct function under the covers
+                            self.assertEquals(patch_direct.call_count, 1 + patch_host.call_count)
+                            self.assertEquals(mock_update_host_plugin_from_goal_state.call_count, 1)
 
-                        self.assertEquals(HostPluginProtocol.is_default_channel(), False)
+                            self.assertEquals(HostPluginProtocol.is_default_channel(), False)
 
-    def test_send_request_using_appropriate_channel_should_not_invoke_host_channel_when_direct_channel_succeeds(self,
-                                                                                                                *args):
-        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
-        client = WireClient(WIRESERVER_URL)
-        client._set_goal_state(GoalState(xml_text))
-        client.get_host_plugin().set_default_channel(False)
-
-        def direct_func(*args):
-            direct_func.counter += 1
-            return 42
-
-        def host_func(*args):
-            host_func.counter += 1
-            return None
-
-        direct_func.counter = 0
-        host_func.counter = 0
-
-        # Assert we've only called the direct channel functions and that it succeeded.
-        ret = client.send_request_using_appropriate_channel(direct_func, host_func)
-        self.assertEquals(42, ret)
-        self.assertEquals(1, direct_func.counter)
-        self.assertEquals(0, host_func.counter)
-
-    def test_send_request_using_appropriate_channel_should_not_use_direct_channel_when_host_channel_is_default(self,
-                                                                                                               *args):
-        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
-        client = WireClient(WIRESERVER_URL)
-        client._set_goal_state(GoalState(xml_text))
-        client.get_host_plugin().set_default_channel(True)
-
-        def direct_func(*args):
-            direct_func.counter += 1
-            return 42
-
-        def host_func(*args):
-            host_func.counter += 1
-            return 43
-
-        direct_func.counter = 0
-        host_func.counter = 0
-
-        # Assert we've only called the host channel function since it's the default channel
-        ret = client.send_request_using_appropriate_channel(direct_func, host_func)
-        self.assertEquals(43, ret)
-        self.assertEquals(0, direct_func.counter)
-        self.assertEquals(1, host_func.counter)
-
-    def test_send_request_using_appropriate_channel_should_use_host_channel_when_direct_channel_fails(self, *args):
-        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
-        client = WireClient(WIRESERVER_URL)
-        client._set_goal_state(GoalState(xml_text))
-        host = client.get_host_plugin()
-        host.set_default_channel(False)
-
-        def direct_func(*args):
-            direct_func.counter += 1
-            raise InvalidContainerError()
-
-        def host_func(*args):
-            host_func.counter += 1
-            return 42
-
-        direct_func.counter = 0
-        host_func.counter = 0
-
-        # Assert we've called both the direct channel function and the host channel function, which succeeded.
-        # After the host channel succeeds, the host plugin should have been set as the default channel.
-        ret = client.send_request_using_appropriate_channel(direct_func, host_func)
-        self.assertEquals(42, ret)
-        self.assertEquals(1, direct_func.counter)
-        self.assertEquals(1, host_func.counter)
-        self.assertEquals(True, host.is_default_channel())
-
-    def test_send_request_using_appropriate_channel_should_retry_the_host_channel_after_reloading_goal_state(self,
-                                                                                                             *args):
-        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
-        client = WireClient(WIRESERVER_URL)
-        client._set_goal_state(GoalState(xml_text))
-        client.get_host_plugin().set_default_channel(False)
-
-        def direct_func(*args):
-            direct_func.counter += 1
-            raise InvalidContainerError()
-
-        def host_func(*args):
-            host_func.counter += 1
-            if host_func.counter == 1:
-                raise ResourceGoneError("Resource is gone")
-            return 42
-
-        direct_func.counter = 0
-        host_func.counter = 0
-
-        # Assert we've called both the direct channel function (once) and the host channel function (twice).
-        # After the host channel succeeds, the host plugin should have been set as the default channel.
-        with patch(
-                'azurelinuxagent.common.protocol.wire.WireClient.update_host_plugin_from_goal_state') as mock_update_host_plugin_from_goal_state:
-            ret = client.send_request_using_appropriate_channel(direct_func, host_func)
+    def test_send_request_using_appropriate_channel_should_not_invoke_host_channel_when_direct_channel_succeeds(self):
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE) as protocol:
+            protocol.client.get_host_plugin().set_default_channel(False)
+
+            def direct_func(*args):
+                direct_func.counter += 1
+                return 42
+
+            def host_func(*args):
+                host_func.counter += 1
+                return None
+
+            direct_func.counter = 0
+            host_func.counter = 0
+
+            # Assert we've only called the direct channel functions and that it succeeded.
+            ret = protocol.client.send_request_using_appropriate_channel(direct_func, host_func)
             self.assertEquals(42, ret)
             self.assertEquals(1, direct_func.counter)
-            self.assertEquals(2, host_func.counter)
-            self.assertEquals(1, mock_update_host_plugin_from_goal_state.call_count)
-            self.assertEquals(True, client.get_host_plugin().is_default_channel())
+            self.assertEquals(0, host_func.counter)
 
+    def test_send_request_using_appropriate_channel_should_not_use_direct_channel_when_host_channel_is_default(self):
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE) as protocol:
+            protocol.client.get_host_plugin().set_default_channel(True)
 
-class GoalStateTestCase(AgentTestCase):
-    """
-    Base class for UpdateGoalStateTestCase and UpdateHostPluginFromGoalStateTestCase
-    """
+            def direct_func(*args):
+                direct_func.counter += 1
+                return 42
 
-    def setUp(self):
-        AgentTestCase.setUp(self)
-
-        self.mock_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
-        self.mock_http_get = patch("azurelinuxagent.common.utils.restutil.http_get",
-                                   side_effect=self.mock_data.mock_http_get)
-        self.mock_crypt_util = patch("azurelinuxagent.common.protocol.wire.CryptUtil",
-                                     side_effect=self.mock_data.mock_crypt_util)
-        self.mock_http_get.start()
-        self.mock_crypt_util.start()
-
-        # Use the WireProtocol to create a test WireClient.
-        # WireProtocol.detect() creates the transport certs needed to retrieve the goal state and also sets the initial goal state.
-        # WireClient.get_host_plugin() forces the host plugin to be instantiated
-        protocol = WireProtocol(WIRESERVER_URL)
-        protocol.detect()
-        self.wire_client = protocol.client
-        self.wire_client.get_host_plugin()
-
-    def tearDown(self):
-        self.mock_crypt_util.stop()
-        self.mock_http_get.stop()
-        AgentTestCase.tearDown(self)
-
-
-class UpdateGoalStateTestCase(GoalStateTestCase):
-    """
-    Tests for WireClient.update_goal_state()
-    """
+            def host_func(*args):
+                host_func.counter += 1
+                return 43
 
-    def test_it_should_update_the_goal_state_and_the_host_plugin_when_the_incarnation_changes(self):
-        # if the incarnation changes the behavior is the same for forced and non-forced updates
-        for forced in [True, False]:
-            self.mock_data.reload()  # start each iteration of the test with fresh mock data
-
-            #
-            # Update the mock data with random values; include at least one field from each of the components
-            # in the goal state to ensure the entire state was updated. Note that numeric entities, e.g. incarnation, are
-            # actually represented as strings in the goal state.
-            #
-            # Note that the shared config is not parsed by the agent, so we modify the XML data directly. Also, the
-            # certificates are encrypted and it is hard to update a single field; instead, we update the entire list with
-            # empty.
-            #
-            new_incarnation = str(uuid.uuid4())
-            new_container_id = str(uuid.uuid4())
-            new_role_config_name = str(uuid.uuid4())
-            new_hosting_env_deployment_name = str(uuid.uuid4())
-            new_shared_conf = WireProtocolData.replace_xml_attribute_value(self.mock_data.shared_config, "Deployment",
-                                                                           "name", str(uuid.uuid4()))
-            new_sequence_number = str(uuid.uuid4())
-
-            if '<Format>Pkcs7BlobWithPfxContents</Format>' not in self.mock_data.certs:
-                raise Exception('This test requires a non-empty certificate list')
-
-            self.mock_data.set_incarnation(new_incarnation)
-            self.mock_data.set_container_id(new_container_id)
-            self.mock_data.set_role_config_name(new_role_config_name)
-            self.mock_data.set_hosting_env_deployment_name(new_hosting_env_deployment_name)
-            self.mock_data.shared_config = new_shared_conf
-            self.mock_data.set_extensions_config_sequence_number(new_sequence_number)
-            self.mock_data.certs = r'''<?xml version="1.0" encoding="utf-8"?>
-                <CertificateFile><Version>2012-11-30</Version>
-                  <Incarnation>12</Incarnation>
-                  <Format>CertificatesNonPfxPackage</Format>
-                  <Data>NotPFXData</Data>
-                </CertificateFile>
-            '''
-
-            if forced:
-                self.wire_client.update_goal_state(forced=True)
-            else:
-                self.wire_client.update_goal_state()
+            direct_func.counter = 0
+            host_func.counter = 0
+
+            # Assert we've only called the host channel function since it's the default channel
+            ret = protocol.client.send_request_using_appropriate_channel(direct_func, host_func)
+            self.assertEquals(43, ret)
+            self.assertEquals(0, direct_func.counter)
+            self.assertEquals(1, host_func.counter)
+
+    def test_send_request_using_appropriate_channel_should_use_host_channel_when_direct_channel_fails(self):
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE) as protocol:
+            host = protocol.client.get_host_plugin()
+            host.set_default_channel(False)
+
+            def direct_func(*args):
+                direct_func.counter += 1
+                raise InvalidContainerError()
+
+            def host_func(*args):
+                host_func.counter += 1
+                return 42
+
+            direct_func.counter = 0
+            host_func.counter = 0
+
+            # Assert we've called both the direct channel function and the host channel function, which succeeded.
+            # After the host channel succeeds, the host plugin should have been set as the default channel.
+            ret = protocol.client.send_request_using_appropriate_channel(direct_func, host_func)
+            self.assertEquals(42, ret)
+            self.assertEquals(1, direct_func.counter)
+            self.assertEquals(1, host_func.counter)
+            self.assertEquals(True, host.is_default_channel())
 
-            sequence_number = self.wire_client.get_ext_conf().ext_handlers.extHandlers[0].properties.extensions[
-                0].sequenceNumber
+    def test_send_request_using_appropriate_channel_should_retry_the_host_channel_after_reloading_goal_state(self):
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE) as protocol:
+            protocol.client.get_host_plugin().set_default_channel(False)
 
-            self.assertEqual(self.wire_client.get_goal_state().incarnation, new_incarnation)
-            self.assertEqual(self.wire_client.get_hosting_env().deployment_name, new_hosting_env_deployment_name)
-            self.assertEqual(self.wire_client.get_shared_conf().xml_text, new_shared_conf)
-            self.assertEqual(sequence_number, new_sequence_number)
-            self.assertEqual(len(self.wire_client.get_certs().cert_list.certificates), 0)
+            def direct_func(*args):
+                direct_func.counter += 1
+                raise InvalidContainerError()
 
-            self.assertEqual(self.wire_client.get_host_plugin().container_id, new_container_id)
-            self.assertEqual(self.wire_client.get_host_plugin().role_config_name, new_role_config_name)
+            def host_func(*args):
+                host_func.counter += 1
+                if host_func.counter == 1:
+                    raise ResourceGoneError("Resource is gone")
+                return 42
 
-    def test_non_forced_update_should_not_update_the_goal_state_nor_the_host_plugin_when_the_incarnation_does_not_change(
-            self):
-        # The container id, role config name and shared config can change without the incarnation changing; capture the initial
-        # goal state and then change those fields.
-        goal_state = self.wire_client.get_goal_state().xml_text
-        shared_conf = self.wire_client.get_shared_conf().xml_text
-        container_id = self.wire_client.get_host_plugin().container_id
-        role_config_name = self.wire_client.get_host_plugin().role_config_name
+            direct_func.counter = 0
+            host_func.counter = 0
 
-        self.mock_data.set_container_id(str(uuid.uuid4()))
-        self.mock_data.set_role_config_name(str(uuid.uuid4()))
-        self.mock_data.shared_config = WireProtocolData.replace_xml_attribute_value(self.mock_data.shared_config,
-                                                                                    "Deployment", "name",
-                                                                                    str(uuid.uuid4()))
+            # Assert we've called both the direct channel function (once) and the host channel function (twice).
+            # After the host channel succeeds, the host plugin should have been set as the default channel.
+            with patch(
+                    'azurelinuxagent.common.protocol.wire.WireClient.update_host_plugin_from_goal_state') as mock_update_host_plugin_from_goal_state:
+                ret = protocol.client.send_request_using_appropriate_channel(direct_func, host_func)
+                self.assertEquals(42, ret)
+                self.assertEquals(1, direct_func.counter)
+                self.assertEquals(2, host_func.counter)
+                self.assertEquals(1, mock_update_host_plugin_from_goal_state.call_count)
+                self.assertEquals(True, protocol.client.get_host_plugin().is_default_channel())
 
-        self.wire_client.update_goal_state()
 
-        self.assertEqual(self.wire_client.get_goal_state().xml_text, goal_state)
-        self.assertEqual(self.wire_client.get_shared_conf().xml_text, shared_conf)
+class UpdateGoalStateTestCase(AgentTestCase):
+    """
+    Tests for WireClient.update_goal_state()
+    """
 
-        self.assertEqual(self.wire_client.get_host_plugin().container_id, container_id)
-        self.assertEqual(self.wire_client.get_host_plugin().role_config_name, role_config_name)
+    def test_it_should_update_the_goal_state_and_the_host_plugin_when_the_incarnation_changes(self):
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE) as protocol:
+            protocol.client.get_host_plugin()
+
+            # if the incarnation changes the behavior is the same for forced and non-forced updates
+            for forced in [True, False]:
+                protocol.mock_wire_data.reload()  # start each iteration of the test with fresh mock data
+
+                #
+                # Update the mock data with random values; include at least one field from each of the components
+                # in the goal state to ensure the entire state was updated. Note that numeric entities, e.g. incarnation, are
+                # actually represented as strings in the goal state.
+                #
+                # Note that the shared config is not parsed by the agent, so we modify the XML data directly. Also, the
+                # certificates are encrypted and it is hard to update a single field; instead, we update the entire list with
+                # empty.
+                #
+                new_incarnation = str(uuid.uuid4())
+                new_container_id = str(uuid.uuid4())
+                new_role_config_name = str(uuid.uuid4())
+                new_hosting_env_deployment_name = str(uuid.uuid4())
+                new_shared_conf = WireProtocolData.replace_xml_attribute_value(protocol.mock_wire_data.shared_config, "Deployment", "name", str(uuid.uuid4()))
+                new_sequence_number = str(uuid.uuid4())
+
+                if '<Format>Pkcs7BlobWithPfxContents</Format>' not in protocol.mock_wire_data.certs:
+                    raise Exception('This test requires a non-empty certificate list')
+
+                protocol.mock_wire_data.set_incarnation(new_incarnation)
+                protocol.mock_wire_data.set_container_id(new_container_id)
+                protocol.mock_wire_data.set_role_config_name(new_role_config_name)
+                protocol.mock_wire_data.set_hosting_env_deployment_name(new_hosting_env_deployment_name)
+                protocol.mock_wire_data.shared_config = new_shared_conf
+                protocol.mock_wire_data.set_extensions_config_sequence_number(new_sequence_number)
+                protocol.mock_wire_data.certs = r'''<?xml version="1.0" encoding="utf-8"?>
+                    <CertificateFile><Version>2012-11-30</Version>
+                      <Incarnation>12</Incarnation>
+                      <Format>CertificatesNonPfxPackage</Format>
+                      <Data>NotPFXData</Data>
+                    </CertificateFile>
+                '''
+
+                if forced:
+                    protocol.client.update_goal_state(forced=True)
+                else:
+                    protocol.client.update_goal_state()
+
+                sequence_number = protocol.client.get_ext_conf().ext_handlers.extHandlers[0].properties.extensions[0].sequenceNumber
+
+                self.assertEqual(protocol.client.get_goal_state().incarnation, new_incarnation)
+                self.assertEqual(protocol.client.get_hosting_env().deployment_name, new_hosting_env_deployment_name)
+                self.assertEqual(protocol.client.get_shared_conf().xml_text, new_shared_conf)
+                self.assertEqual(sequence_number, new_sequence_number)
+                self.assertEqual(len(protocol.client.get_certs().cert_list.certificates), 0)
+
+                self.assertEqual(protocol.client.get_host_plugin().container_id, new_container_id)
+                self.assertEqual(protocol.client.get_host_plugin().role_config_name, new_role_config_name)
+
+    def test_non_forced_update_should_not_update_the_goal_state_nor_the_host_plugin_when_the_incarnation_does_not_change(self):
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE) as protocol:
+            protocol.client.get_host_plugin()
+
+            # The container id, role config name and shared config can change without the incarnation changing; capture the initial
+            # goal state and then change those fields.
+            goal_state = protocol.client.get_goal_state().xml_text
+            shared_conf = protocol.client.get_shared_conf().xml_text
+            container_id = protocol.client.get_host_plugin().container_id
+            role_config_name = protocol.client.get_host_plugin().role_config_name
+
+            protocol.mock_wire_data.set_container_id(str(uuid.uuid4()))
+            protocol.mock_wire_data.set_role_config_name(str(uuid.uuid4()))
+            protocol.mock_wire_data.shared_config = WireProtocolData.replace_xml_attribute_value(
+                protocol.mock_wire_data.shared_config, "Deployment", "name", str(uuid.uuid4()))
+
+            protocol.client.update_goal_state()
+
+            self.assertEqual(protocol.client.get_goal_state().xml_text, goal_state)
+            self.assertEqual(protocol.client.get_shared_conf().xml_text, shared_conf)
+
+            self.assertEqual(protocol.client.get_host_plugin().container_id, container_id)
+            self.assertEqual(protocol.client.get_host_plugin().role_config_name, role_config_name)
 
     def test_forced_update_should_update_the_goal_state_and_the_host_plugin_when_the_incarnation_does_not_change(self):
-        # The container id, role config name and shared config can change without the incarnation changing
-        incarnation = self.wire_client.get_goal_state().incarnation
-        new_container_id = str(uuid.uuid4())
-        new_role_config_name = str(uuid.uuid4())
-        new_shared_conf = WireProtocolData.replace_xml_attribute_value(self.mock_data.shared_config, "Deployment",
-                                                                       "name", str(uuid.uuid4()))
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE) as protocol:
+            protocol.client.get_host_plugin()
+
+            # The container id, role config name and shared config can change without the incarnation changing
+            incarnation = protocol.client.get_goal_state().incarnation
+            new_container_id = str(uuid.uuid4())
+            new_role_config_name = str(uuid.uuid4())
+            new_shared_conf = WireProtocolData.replace_xml_attribute_value(
+                protocol.mock_wire_data.shared_config, "Deployment", "name", str(uuid.uuid4()))
 
-        self.mock_data.set_container_id(new_container_id)
-        self.mock_data.set_role_config_name(new_role_config_name)
-        self.mock_data.shared_config = new_shared_conf
+            protocol.mock_wire_data.set_container_id(new_container_id)
+            protocol.mock_wire_data.set_role_config_name(new_role_config_name)
+            protocol.mock_wire_data.shared_config = new_shared_conf
 
-        self.wire_client.update_goal_state(forced=True)
+            protocol.client.update_goal_state(forced=True)
 
-        self.assertEqual(self.wire_client.get_goal_state().incarnation, incarnation)
-        self.assertEqual(self.wire_client.get_shared_conf().xml_text, new_shared_conf)
+            self.assertEqual(protocol.client.get_goal_state().incarnation, incarnation)
+            self.assertEqual(protocol.client.get_shared_conf().xml_text, new_shared_conf)
 
-        self.assertEqual(self.wire_client.get_host_plugin().container_id, new_container_id)
-        self.assertEqual(self.wire_client.get_host_plugin().role_config_name, new_role_config_name)
+            self.assertEqual(protocol.client.get_host_plugin().container_id, new_container_id)
+            self.assertEqual(protocol.client.get_host_plugin().role_config_name, new_role_config_name)
 
 
-class UpdateHostPluginFromGoalStateTestCase(GoalStateTestCase):
+class UpdateHostPluginFromGoalStateTestCase(AgentTestCase):
     """
     Tests for WireClient.update_host_plugin_from_goal_state()
     """
 
     def test_it_should_update_the_host_plugin_with_or_without_incarnation_changes(self):
-        # the behavior should be the same whether the incarnation changes or not
-        for incarnation_change in [True, False]:
-            self.mock_data.reload()  # start each iteration of the test with fresh mock data
+        with mock_wire_protocol.create(mockwiredata.DATA_FILE) as protocol:
+            protocol.client.get_host_plugin()
 
-            new_container_id = str(uuid.uuid4())
-            new_role_config_name = str(uuid.uuid4())
+            # the behavior should be the same whether the incarnation changes or not
+            for incarnation_change in [True, False]:
+                protocol.mock_wire_data.reload()  # start each iteration of the test with fresh mock data
+
+                new_container_id = str(uuid.uuid4())
+                new_role_config_name = str(uuid.uuid4())
+
+                goal_state_xml_text = protocol.mock_wire_data.goal_state
+                shared_conf_xml_text = protocol.mock_wire_data.shared_config
+
+                if incarnation_change:
+                    protocol.mock_wire_data.set_incarnation(str(uuid.uuid4()))
 
-            goal_state = self.mock_data.goal_state
-            shared_conf = self.mock_data.shared_config
-            if incarnation_change:
-                self.mock_data.set_incarnation(str(uuid.uuid4()))
-            self.mock_data.set_container_id(new_container_id)
-            self.mock_data.set_role_config_name(new_role_config_name)
-            self.mock_data.shared_config = WireProtocolData.replace_xml_attribute_value(self.mock_data.shared_config,
-                                                                                        "Deployment", "name",
-                                                                                        str(uuid.uuid4()))
+                protocol.mock_wire_data.set_container_id(new_container_id)
+                protocol.mock_wire_data.set_role_config_name(new_role_config_name)
+                protocol.mock_wire_data.shared_config = WireProtocolData.replace_xml_attribute_value(
+                    protocol.mock_wire_data.shared_config, "Deployment", "name", str(uuid.uuid4()))
 
-            self.wire_client.update_host_plugin_from_goal_state()
+                protocol.client.update_host_plugin_from_goal_state()
 
-            self.assertEqual(self.wire_client.get_host_plugin().container_id, new_container_id)
-            self.assertEqual(self.wire_client.get_host_plugin().role_config_name, new_role_config_name)
+                self.assertEqual(protocol.client.get_host_plugin().container_id, new_container_id)
+                self.assertEqual(protocol.client.get_host_plugin().role_config_name, new_role_config_name)
 
-            # it should not update the goal state
-            self.assertEqual(self.wire_client.get_goal_state().xml_text, goal_state)
-            self.assertEqual(self.wire_client.get_shared_conf().xml_text, shared_conf)
+                # it should not update the goal state
+                self.assertEqual(protocol.client.get_goal_state().xml_text, goal_state_xml_text)
+                self.assertEqual(protocol.client.get_shared_conf().xml_text, shared_conf_xml_text)
 
 
 class MockResponse: