Repository: Azure/WALinuxAgent
Pull request #1741 (2019-12-21 00:36:47)

Do not update goal state when refreshing the host plugin
We are updating the goal state too frequently and that is part of the reason different components are getting out of sync. One of the reasons we were forcing updates of the goal state was to retry calls to the host plugin when the container id or the role config change. In this case updating the full goal state is not needed and it can actually introduce consistency issues.

The main change in this PR is to define an API specifically to update the host plugin without changing the goal state.

Other small changes are minor cleanup of the update goal state in the metadata protocol and the parse code of the goal state; I'll add comments in the PR to explain those changes.

<!-- 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/1741)
<!-- Reviewable:end -->


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

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

```diff
@@            Coverage Diff             @@
##             develop    #1741   +/-   ##
==========================================
  Coverage           ?   68.44%           
==========================================
  Files              ?       81           
  Lines              ?    11734           
  Branches           ?     1645           
==========================================
  Hits               ?     8031           
  Misses             ?     3362           
  Partials           ?      341
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/1741?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/protocol/metadata.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1741/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC9tZXRhZGF0YS5weQ==) | `53.37% <100%> (ø)` | |
| [azurelinuxagent/daemon/main.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1741/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2RhZW1vbi9tYWluLnB5) | `73.68% <33.33%> (ø)` | |
| [azurelinuxagent/common/protocol/wire.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1741/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC93aXJlLnB5) | `76.65% <92.64%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1741?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/1741?src=pr&el=footer). Last update [8194bbc...e1176a6](https://codecov.io/gh/Azure/WALinuxAgent/pull/1741?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/azurelinuxagent/common/protocol/metadata.py b/azurelinuxagent/common/protocol/metadata.py
index bb2e1649..69fd71cc 100644
--- a/azurelinuxagent/common/protocol/metadata.py
+++ b/azurelinuxagent/common/protocol/metadata.py
@@ -170,7 +170,7 @@ class MetadataProtocol(Protocol):
                                 "{0}.crt".format(thumbprint))
         shutil.copyfile(trans_prv_file, prv_file)
         shutil.copyfile(trans_cert_file, crt_file)
-        self.update_goal_state(forced=True)
+        self.update_goal_state()
 
     def get_vminfo(self):
         vminfo = VMInfo()
@@ -330,7 +330,7 @@ class MetadataProtocol(Protocol):
         certificates = self.get_certs()
         return certificates.cert_list
 
-    def update_goal_state(self, forced=False, max_retry=3):
+    def update_goal_state(self, max_retry=3):
         # Start updating goalstate, retry on 410
         for retry in range(0, max_retry):
             try:
diff --git a/azurelinuxagent/common/protocol/wire.py b/azurelinuxagent/common/protocol/wire.py
index 93c30900..2925fe3d 100644
--- a/azurelinuxagent/common/protocol/wire.py
+++ b/azurelinuxagent/common/protocol/wire.py
@@ -106,10 +106,8 @@ class WireProtocol(Protocol):
         cryptutil = CryptUtil(conf.get_openssl_cmd())
         cryptutil.gen_transport_cert(trans_prv_file, trans_cert_file)
 
-        self.update_goal_state(forced=True)
-
-    def update_goal_state(self, forced=False, max_retry=3):
-        self.client.update_goal_state(forced=forced, max_retry=max_retry)
+        # Set the initial goal state
+        self.client.update_goal_state(forced=True)
 
     def get_vminfo(self):
         goal_state = self.client.get_goal_state()
@@ -136,7 +134,7 @@ class WireProtocol(Protocol):
 
     def get_vmagent_manifests(self):
         # Update goal state to get latest extensions config
-        self.update_goal_state()
+        self.client.update_goal_state()
         goal_state = self.client.get_goal_state()
         ext_conf = self.client.get_ext_conf()
         return ext_conf.vmagent_manifests, goal_state.incarnation
@@ -150,7 +148,7 @@ class WireProtocol(Protocol):
     def get_ext_handlers(self):
         logger.verbose("Get extension handler config")
         # Update goal state to get latest extensions config
-        self.update_goal_state()
+        self.client.update_goal_state()
         goal_state = self.client.get_goal_state()
         ext_conf = self.client.get_ext_conf()
         # In wire protocol, incarnation is equivalent to ETag
@@ -761,78 +759,88 @@ class WireClient(object):
         self.save_cache(local_file, xml_text)
         self.ext_conf = ExtensionsConfig(xml_text)
 
-    def save_or_update_goal_state_file(self, incarnation, xml_text):
-        # It should create a new file if the incarnation number is new.
-        # It should overwrite the existing file if the incarnation number is the same.
-        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)
+    # Type of update performed by _update_from_goal_state()
+    class _UpdateType(object):
+        # Update the Host GA Plugin client (Container ID and RoleConfigName)
+        HostPlugin = 0
+        # Update the full goal state only if the incarnation has changed
+        GoalState = 1
+        # Update the full goal state unconditionally
+        GoalStateForced = 2
+
+    def update_host_plugin_from_goal_state(self):
+        """
+        Fetches a new goal state and updates the Container ID and Role Config Name of the host plugin client
+        """
+        self._update_from_goal_state(WireClient._UpdateType.HostPlugin)
 
-    def update_host_plugin(self, container_id, role_config_name):
-        if self.host_plugin is not None:
-            self.host_plugin.container_id = container_id
-            self.host_plugin.role_config_name = role_config_name
+    def update_goal_state(self, forced=False):
+        """
+        Updates the goal state if the incarnation changed or if 'forced' is True
+        """
+        self._update_from_goal_state(WireClient._UpdateType.GoalStateForced if forced else WireClient._UpdateType.GoalState)
 
-    def update_goal_state(self, forced=False, max_retry=3):
+    def _update_from_goal_state(self, refresh_type):
+        """
+        Fetches a new goal state and updates the internal state of the WireClient according to the requested 'refresh_type'
+        """
         incarnation_file = os.path.join(conf.get_lib_dir(), INCARNATION_FILE_NAME)
         uri = GOAL_STATE_URI.format(self.endpoint)
 
-        current_goal_state_from_configuration = None
-        for retry in range(0, max_retry):
+        max_retry = 3
+
+        for retry in range(1, max_retry + 1):
             try:
-                if current_goal_state_from_configuration is None:
-                    xml_text = self.fetch_config(uri, self.get_header())
-                    current_goal_state_from_configuration = GoalState(xml_text)
-
-                    if not forced:
-                        last_incarnation = None
-                        if os.path.isfile(incarnation_file):
-                            last_incarnation = fileutil.read_file(incarnation_file)
-                        new_incarnation = current_goal_state_from_configuration.incarnation
-
-                        if last_incarnation is not None and last_incarnation == new_incarnation:
-                            # Incarnation number is not updated, but role config file and container ID
-                            # can change without the incarnation number changing. Ensure they are updated in
-                            # the goal state file on disk, as well as in the HostGA plugin instance.
-                            self.goal_state = current_goal_state_from_configuration
-                            self.save_or_update_goal_state_file(new_incarnation, xml_text)
-                            self.update_host_plugin(current_goal_state_from_configuration.container_id,
-                                                    current_goal_state_from_configuration.role_config_name)
-
-                            return
-                self.goal_state_flusher.flush(datetime.utcnow())
-
-                self.goal_state = current_goal_state_from_configuration
-                self.save_or_update_goal_state_file(current_goal_state_from_configuration.incarnation, xml_text)
-                self.update_hosting_env(current_goal_state_from_configuration)
-                self.update_shared_conf(current_goal_state_from_configuration)
-                self.update_certs(current_goal_state_from_configuration)
-                self.update_ext_conf(current_goal_state_from_configuration)
-                self.update_remote_access_conf(current_goal_state_from_configuration)
-                self.save_cache(incarnation_file, current_goal_state_from_configuration.incarnation)
-                self.update_host_plugin(current_goal_state_from_configuration.container_id,
-                                        current_goal_state_from_configuration.role_config_name)
+                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.container_id = new_goal_state.container_id
+                        self.host_plugin.role_config_name = new_goal_state.role_config_name
+
+                if refresh_type == WireClient._UpdateType.HostPlugin:
+                    update_host_plugin()
+                    return
+
+                def incarnation_changed():
+                    last_incarnation = None
+                    if os.path.isfile(incarnation_file):
+                        last_incarnation = fileutil.read_file(incarnation_file)
+                    return last_incarnation is None or last_incarnation != new_goal_state.incarnation
+
+                if refresh_type == WireClient._UpdateType.GoalStateForced or incarnation_changed() or self.goal_state is None:
+                    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.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(incarnation_file, new_goal_state.incarnation)
+
+                    update_host_plugin()
 
                 return
 
             except IOError as e:
-                logger.warn("IOError processing goal state, retrying [{0}]", ustr(e))
+                logger.warn("IOError processing goal state (attempt {0}/{1}) [{2}]", retry, max_retry, ustr(e))
 
             except ResourceGoneError:
-                logger.info("Goal state is stale, re-fetching")
-                current_goal_state_from_configuration = None
+                logger.info("Goal state is stale, re-fetching (attempt {0}/{1})", retry, max_retry)
 
             except ProtocolError as e:
-                if retry < max_retry - 1:
-                    logger.verbose("ProtocolError processing goal state, retrying [{0}]", ustr(e))
-                else:
-                    logger.error("ProtocolError processing goal state, giving up [{0}]", ustr(e))
+                logger.verbose("ProtocolError processing goal state (attempt {0}/{1}) [{2}]", retry, max_retry, ustr(e))
 
             except Exception as e:
-                if retry < max_retry - 1:
-                    logger.verbose("Exception processing goal state, retrying: [{0}]", ustr(e))
-                else:
-                    logger.error("Exception processing goal state, giving up: [{0}]", ustr(e))
+                logger.verbose("Exception processing goal state (attempt {0}/{1}) [{2}]", retry, max_retry, ustr(e))
 
         raise ProtocolError("Exceeded max retry updating goal state")
 
@@ -1041,7 +1049,7 @@ class WireClient(object):
                   "Error: {2}".format(old_container_id, old_role_config_name, ustr(e))
             logger.periodic_info(logger.EVERY_SIX_HOURS, msg)
 
-            self.update_goal_state(forced=True)
+            self.update_host_plugin_from_goal_state()
 
             new_container_id = self.host_plugin.container_id
             new_role_config_name = self.host_plugin.role_config_name
@@ -1090,6 +1098,7 @@ class WireClient(object):
         ext_conf = self.get_ext_conf()
 
         if ext_conf.status_upload_blob is None:
+            # the status upload blob is in ExtensionsConfig so force a full goal state refresh
             self.update_goal_state(forced=True)
             ext_conf = self.get_ext_conf()
 
@@ -1120,8 +1129,8 @@ class WireClient(object):
             host.put_vm_status(self.status_blob, ext_conf.status_upload_blob, ext_conf.status_upload_blob_type)
             return
         except ResourceGoneError:
-            # do not attempt direct, force goal state update and wait to try again
-            self.update_goal_state(forced=True)
+            # refresh the host plugin client and try again on the next iteration of the main loop
+            self.update_host_plugin_from_goal_state()
             return
         except Exception as e:
             # for all other errors, fall back to direct
@@ -1369,12 +1378,9 @@ class GoalState(object):
         self.container_id = None
         self.load_balancer_probe_port = None
         self.xml_text = None
-        self.parse(xml_text)
+        self._parse(xml_text)
 
-    def parse(self, xml_text):
-        """
-        Request configuration data from endpoint server.
-        """
+    def _parse(self, xml_text):
         self.xml_text = xml_text
         xml_doc = parse_doc(xml_text)
         self.incarnation = findtext(xml_doc, "Incarnation")
@@ -1393,15 +1399,9 @@ class GoalState(object):
         self.remote_access_uri = findtext(container, "RemoteAccessInfo")
         lbprobe_ports = find(xml_doc, "LBProbePorts")
         self.load_balancer_probe_port = findtext(lbprobe_ports, "Port")
-        return self
 
 
 class HostingEnv(object):
-    """
-    parse Hosting enviromnet config and store in
-    HostingEnvironmentConfig.xml
-    """
-
     def __init__(self, xml_text):
         if xml_text is None:
             raise ValueError("HostingEnvironmentConfig.xml is None")
@@ -1410,12 +1410,9 @@ class HostingEnv(object):
         self.role_name = None
         self.deployment_name = None
         self.xml_text = None
-        self.parse(xml_text)
+        self._parse(xml_text)
 
-    def parse(self, xml_text):
-        """
-        parse and create HostingEnvironmentConfig.xml.
-        """
+    def _parse(self, xml_text):
         self.xml_text = xml_text
         xml_doc = parse_doc(xml_text)
         incarnation = find(xml_doc, "Incarnation")
@@ -1424,24 +1421,12 @@ class HostingEnv(object):
         self.role_name = getattrib(role, "name")
         deployment = find(xml_doc, "Deployment")
         self.deployment_name = getattrib(deployment, "name")
-        return self
 
 
 class SharedConfig(object):
-    """
-    parse role endpoint server and goal state config.
-    """
-
     def __init__(self, xml_text):
         logger.verbose("Load SharedConfig.xml")
-        self.parse(xml_text)
-
-    def parse(self, xml_text):
-        """
-        parse and write configuration to file SharedConfig.xml.
-        """
-        # Not used currently
-        return self
+        self.xml_text = xml_text
 
 
 class RemoteAccess(object):
@@ -1470,9 +1455,9 @@ class RemoteAccess(object):
         self.user_list = RemoteAccessUsersList()
 
         self.xml_text = None
-        self.parse(xml_text)
+        self._parse(xml_text)
 
-    def parse(self, xml_text):
+    def _parse(self, xml_text):
         """
         Parse xml document containing user account information
         """
@@ -1486,11 +1471,10 @@ class RemoteAccess(object):
         users = findall(user_collection, "User")
 
         for user in users:
-            remote_access_user = self.parse_user(user)
+            remote_access_user = self._parse_user(user)
             self.user_list.users.append(remote_access_user)
-        return self
 
-    def parse_user(self, user):
+    def _parse_user(self, user):
         name = findtext(user, "Name")
         encrypted_password = findtext(user, "Password")
         expiration = findtext(user, "Expiration")
@@ -1520,9 +1504,9 @@ class Certificates(object):
         logger.verbose("Load Certificates.xml")
         self.client = client
         self.cert_list = CertList()
-        self.parse(xml_text)
+        self._parse(xml_text)
 
-    def parse(self, xml_text):
+    def _parse(self, xml_text):
         """
         Parse multiple certificates into seperate files.
         """
@@ -1628,11 +1612,6 @@ class Certificates(object):
 
 
 class ExtensionsConfig(object):
-    """
-    parse ExtensionsConfig, downloading and unpacking them to /var/lib/waagent.
-    Install if <enabled>true</enabled>, remove if it is set to false.
-    """
-
     def __init__(self, xml_text):
         logger.verbose("Load ExtensionsConfig.xml")
         self.ext_handlers = ExtHandlerList()
@@ -1641,12 +1620,9 @@ class ExtensionsConfig(object):
         self.status_upload_blob_type = None
         self.artifacts_profile_blob = None
         if xml_text is not None:
-            self.parse(xml_text)
+            self._parse(xml_text)
 
-    def parse(self, xml_text):
-        """
-        Write configuration to file ExtensionsConfig.xml.
-        """
+    def _parse(self, xml_text):
         xml_doc = parse_doc(xml_text)
 
         ga_families_list = find(xml_doc, "GAFamilies")
@@ -1669,9 +1645,9 @@ class ExtensionsConfig(object):
         plugin_settings = findall(plugin_settings_list, "Plugin")
 
         for plugin in plugins:
-            ext_handler = self.parse_plugin(plugin)
+            ext_handler = self._parse_plugin(plugin)
             self.ext_handlers.extHandlers.append(ext_handler)
-            self.parse_plugin_settings(ext_handler, plugin_settings)
+            self._parse_plugin_settings(ext_handler, plugin_settings)
 
         self.status_upload_blob = findtext(xml_doc, "StatusUploadBlob")
         self.artifacts_profile_blob = findtext(xml_doc, "InVMArtifactsProfileBlob")
@@ -1682,7 +1658,7 @@ class ExtensionsConfig(object):
         logger.verbose("Extension config shows status blob type as [{0}]",
                        self.status_upload_blob_type)
 
-    def parse_plugin(self, plugin):
+    def _parse_plugin(self, plugin):
         ext_handler = ExtHandler()
         ext_handler.name = getattrib(plugin, "name")
         ext_handler.properties.version = getattrib(plugin, "version")
@@ -1696,7 +1672,7 @@ class ExtensionsConfig(object):
             ext_handler.versionUris.append(version_uri)
         return ext_handler
 
-    def parse_plugin_settings(self, ext_handler, plugin_settings):
+    def _parse_plugin_settings(self, ext_handler, plugin_settings):
         if plugin_settings is None:
             return
 
@@ -1751,9 +1727,9 @@ class ExtensionManifest(object):
         logger.verbose("Load ExtensionManifest.xml")
         self.pkg_list = ExtHandlerPackageList()
         self.allowed_versions = None
-        self.parse(xml_text)
+        self._parse(xml_text)
 
-    def parse(self, xml_text):
+    def _parse(self, xml_text):
         xml_doc = parse_doc(xml_text)
         self._handle_packages(findall(find(xml_doc,
                                            "Plugins"),
diff --git a/azurelinuxagent/daemon/main.py b/azurelinuxagent/daemon/main.py
index a951e394..18ceee96 100644
--- a/azurelinuxagent/daemon/main.py
+++ b/azurelinuxagent/daemon/main.py
@@ -31,7 +31,7 @@ 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.wire import WireClient
+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, \
     AGENT_VERSION, \
@@ -150,10 +150,9 @@ class DaemonHandler(object):
                 #   incarnation number. A forced update ensures the most
                 #   current values.
                 protocol = self.protocol_util.get_protocol()
-                client = protocol.client
-                if client is None or type(client) is not WireClient:
+                if type(protocol) is not WireProtocol:
                     raise Exception("Attempt to setup RDMA without Wireserver")
-                client.update_goal_state(forced=True)
+                protocol.client.update_goal_state(forced=True)
 
                 setup_rdma_device(nd_version)
             except Exception as e:


Test patch:
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index 38cb8712..552c963f 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -1667,7 +1667,6 @@ class ProtocolMock(object):
             "update_goal_state" : 0
         }
         self.goal_state_is_stale = False
-        self.goal_state_forced = False
         self.etag = etag
         self.versions = versions if versions is not None else []
         self.create_manifests()
@@ -1717,9 +1716,8 @@ class ProtocolMock(object):
             raise ResourceGoneError()
         return self.agent_packages
 
-    def update_goal_state(self, forced=False, max_retry=3):
+    def update_goal_state(self, forced=False):
         self.call_counts["update_goal_state"] += 1
-        self.goal_state_forced = self.goal_state_forced or forced
 
 
 class ResponseMock(Mock):
diff --git a/tests/protocol/mockwiredata.py b/tests/protocol/mockwiredata.py
index 5202003e..bca9b415 100644
--- a/tests/protocol/mockwiredata.py
+++ b/tests/protocol/mockwiredata.py
@@ -15,6 +15,8 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
+import re
+
 from tests.tools import load_bin_data, load_data, MagicMock, Mock
 from azurelinuxagent.common.exception import HttpError, ResourceGoneError
 from azurelinuxagent.common.future import httpclient
@@ -28,7 +30,7 @@ DATA_FILE = {
         "certs": "wire/certs.xml",
         "ext_conf": "wire/ext_conf.xml",
         "manifest": "wire/manifest.xml",
-        "ga_manifest" : "wire/ga_manifest.xml",
+        "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"
@@ -77,29 +79,45 @@ class WireProtocolData(object):
     def __init__(self, data_files=DATA_FILE):
         self.emulate_stale_goal_state = False
         self.call_counts = {
-            "comp=versions" : 0,
-            "/versions" : 0,
-            "goalstate" : 0,
-            "hostingenvuri" : 0,
-            "sharedconfiguri" : 0,
-            "certificatesuri" : 0,
-            "extensionsconfiguri" : 0,
-            "extensionArtifact" : 0,
-            "manifest.xml" : 0,
-            "manifest_of_ga.xml" : 0,
-            "ExampleHandlerLinux" : 0
+            "comp=versions": 0,
+            "/versions": 0,
+            "goalstate": 0,
+            "hostingenvuri": 0,
+            "sharedconfiguri": 0,
+            "certificatesuri": 0,
+            "extensionsconfiguri": 0,
+            "extensionArtifact": 0,
+            "manifest.xml": 0,
+            "manifest_of_ga.xml": 0,
+            "ExampleHandlerLinux": 0
         }
-        self.version_info = load_data(data_files.get("version_info"))
-        self.goal_state = load_data(data_files.get("goal_state"))
-        self.hosting_env = load_data(data_files.get("hosting_env"))
-        self.shared_config = load_data(data_files.get("shared_config"))
-        self.certs = load_data(data_files.get("certs"))
-        self.ext_conf = load_data(data_files.get("ext_conf"))
-        self.manifest = load_data(data_files.get("manifest"))
-        self.ga_manifest = load_data(data_files.get("ga_manifest"))
-        self.trans_prv = load_data(data_files.get("trans_prv"))
-        self.trans_cert = load_data(data_files.get("trans_cert"))
-        self.ext = load_bin_data(data_files.get("test_ext"))
+        self.data_files = data_files
+        self.version_info = None
+        self.goal_state = None
+        self.hosting_env = None
+        self.shared_config = None
+        self.certs = None
+        self.ext_conf = None
+        self.manifest = None
+        self.ga_manifest = None
+        self.trans_prv = None
+        self.trans_cert = None
+        self.ext = None
+
+        self.reload()
+
+    def reload(self):
+        self.version_info = load_data(self.data_files.get("version_info"))
+        self.goal_state = load_data(self.data_files.get("goal_state"))
+        self.hosting_env = load_data(self.data_files.get("hosting_env"))
+        self.shared_config = load_data(self.data_files.get("shared_config"))
+        self.certs = load_data(self.data_files.get("certs"))
+        self.ext_conf = load_data(self.data_files.get("ext_conf"))
+        self.manifest = load_data(self.data_files.get("manifest"))
+        self.ga_manifest = load_data(self.data_files.get("ga_manifest"))
+        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"))
 
     def mock_http_get(self, url, *args, **kwargs):
         content = None
@@ -182,3 +200,52 @@ class WireProtocolData(object):
 
         with open(trans_cert_file, 'w+') as cert_file:
             cert_file.write(self.trans_cert)
+
+    #
+    # Having trouble reading the regular expressions below? you are not alone!
+    #
+    # For the use of "(?<=" "(?=" see 7.2.1 in https://docs.python.org/3.1/library/re.html
+    # For the use of "\g<1>" see backreferences in https://docs.python.org/3.1/library/re.html#re.sub
+    #
+    # Note that these regular expressions are not enough to parse all valid XML documents (e.g. they do
+    # not account for metacharacters like < or > in the values) but they are good enough for the test
+    # data. There are some basic checks, but the functions may not match valid XML or produce invalid
+    # XML if their input is too complex.
+    #
+    @staticmethod
+    def replace_xml_element_value(xml_document, element_name, element_value):
+        new_xml_document = re.sub(r'(?<=<{0}>).+(?=</{0}>)'.format(element_name), element_value, xml_document)
+        if new_xml_document == xml_document:
+            raise Exception("Could not match element '{0}'", element_name)
+        return new_xml_document
+
+    @staticmethod
+    def replace_xml_attribute_value(xml_document, element_name, attribute_name, attribute_value):
+        new_xml_document = re.sub(r'(?<=<{0} )(.*{1}=")[^"]+(?="[^>]*>)'.format(element_name, attribute_name), r'\g<1>{0}'.format(attribute_value), xml_document)
+        if new_xml_document == xml_document:
+            raise Exception("Could not match attribute '{0}' of element '{1}'", attribute_name, element_name)
+        return new_xml_document
+
+    def set_incarnation(self, incarnation):
+        '''
+        Sets the incarnation in the goal state, but not on its subcomponents (e.g. hosting env, shared config)
+        '''
+        self.goal_state = WireProtocolData.replace_xml_element_value(self.goal_state, "Incarnation", incarnation)
+
+    def set_container_id(self, container_id):
+        self.goal_state = WireProtocolData.replace_xml_element_value(self.goal_state, "ContainerId", container_id)
+
+    def set_role_config_name(self, role_config_name):
+        self.goal_state = WireProtocolData.replace_xml_element_value(self.goal_state, "ConfigName", role_config_name)
+
+    def set_hosting_env_deployment_name(self, deployment_name):
+        self.hosting_env = WireProtocolData.replace_xml_attribute_value(self.hosting_env, "Deployment", "name", deployment_name)
+
+    def set_shared_config_deployment_name(self, deployment_name):
+        self.shared_config = WireProtocolData.replace_xml_attribute_value(self.shared_config, "Deployment", "name", deployment_name)
+
+    def set_extensions_config_sequence_number(self, sequence_number):
+        '''
+        Sets the sequence number for *all* extensions
+        '''
+        self.ext_conf = WireProtocolData.replace_xml_attribute_value(self.ext_conf, "RuntimeSettings", "seqNo", sequence_number)
diff --git a/tests/protocol/test_hostplugin.py b/tests/protocol/test_hostplugin.py
index 3aa92669..d19eadb5 100644
--- a/tests/protocol/test_hostplugin.py
+++ b/tests/protocol/test_hostplugin.py
@@ -305,7 +305,8 @@ class TestHostPlugin(AgentTestCase):
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol._put_page_blob_status",
            side_effect=ResourceGoneError("410"))
     @patch("azurelinuxagent.common.protocol.wire.WireClient.update_goal_state")
-    def test_fallback_channel_410(self, patch_update, patch_put, patch_upload, _):
+    @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, _):
         """
         When host plugin returns a 410, we should force the goal state update and return
         """
@@ -328,10 +329,10 @@ class TestHostPlugin(AgentTestCase):
         # 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(2, patch_update.call_count, "Update goal state unexpected call count")
-        self.assertEqual(1, len(patch_update.call_args[1]), "Update goal state unexpected call count")
-        self.assertTrue(patch_update.call_args[1]['forced'], "Update goal state 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, 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")
 
         # ensure the correct url is used
         self.assertEqual(sas_url, patch_put.call_args[0][0])
diff --git a/tests/protocol/test_wire.py b/tests/protocol/test_wire.py
index ca162ba4..02b85b01 100644
--- a/tests/protocol/test_wire.py
+++ b/tests/protocol/test_wire.py
@@ -21,10 +21,10 @@ import os
 import stat
 import tempfile
 import time
+import uuid
 import unittest
 import zipfile
 
-from azurelinuxagent.common import conf
 from azurelinuxagent.common.exception import InvalidContainerError, ResourceGoneError, ProtocolError, \
     ExtensionDownloadError
 from azurelinuxagent.common.future import httpclient
@@ -38,6 +38,7 @@ from azurelinuxagent.common.utils.shellutil import run_get_output
 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.tools import ANY, MagicMock, Mock, patch, AgentTestCase, skip_if_predicate_true, running_under_travis
 
 data_with_bom = b'\xef\xbb\xbfhehe'
@@ -140,7 +141,7 @@ class TestWireProtocol(AgentTestCase):
         # -- Tracking calls to retrieve GoalState is problematic since it is
         #    fetched often; however, the dependent documents, such as the
         #    HostingEnvironmentConfig, will be retrieved the expected number
-        self.assertEqual(2, test_data.call_counts["hostingenvuri"])
+        self.assertEqual(1, test_data.call_counts["hostingenvuri"])
         self.assertEqual(1, patch_report.call_count)
 
     def test_call_storage_kwargs(self, *args):
@@ -610,73 +611,6 @@ class TestWireClient(AgentTestCase):
         self.assertEqual("BlockBlob", ext_conf.status_upload_blob_type)
         self.assertEqual(None, ext_conf.artifacts_profile_blob)
 
-    def test_save_or_update_goal_state_should_save_new_goal_state_file(self):
-        # Assert the file didn't exist before
-        incarnation = 42
-        goal_state_file = os.path.join(conf.get_lib_dir(), "GoalState.{0}.xml".format(incarnation))
-        self.assertFalse(os.path.exists(goal_state_file))
-
-        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
-        client = WireClient(WIRESERVER_URL)
-        client.save_or_update_goal_state_file(incarnation, xml_text)
-
-        # Assert the file exists and its contents
-        self.assertTrue(os.path.exists(goal_state_file))
-        with open(goal_state_file, "r") as f:
-            contents = f.readlines()
-            self.assertEquals("".join(contents), xml_text)
-
-    def test_save_or_update_goal_state_should_update_existing_goal_state_file(self):
-        incarnation = 42
-        goal_state_file = os.path.join(conf.get_lib_dir(), "GoalState.{0}.xml".format(incarnation))
-        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
-
-        with open(goal_state_file, "w") as f:
-            f.write(xml_text)
-
-        # Assert the file exists and its contents
-        self.assertTrue(os.path.exists(goal_state_file))
-        with open(goal_state_file, "r") as f:
-            contents = f.readlines()
-            self.assertEquals("".join(contents), xml_text)
-
-        # Update the container id
-        new_goal_state = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state.replace("c6d5526c-5ac2-4200-b6e2-56f2b70c5ab2",
-                                                                        "z6d5526c-5ac2-4200-b6e2-56f2b70c5ab2")
-        client = WireClient(WIRESERVER_URL)
-        client.save_or_update_goal_state_file(incarnation, new_goal_state)
-
-        # Assert the file exists and its contents
-        self.assertTrue(os.path.exists(goal_state_file))
-        with open(goal_state_file, "r") as f:
-            contents = f.readlines()
-            self.assertEquals("".join(contents), new_goal_state)
-
-    def test_save_or_update_goal_state_should_update_goal_state_and_container_id_when_not_forced(self):
-        incarnation = "1"  # Match the incarnation number from dummy goal state file
-        incarnation_file = os.path.join(conf.get_lib_dir(), INCARNATION_FILE_NAME)
-        with open(incarnation_file, "w") as f:
-            f.write(incarnation)
-
-        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
-        goal_state_file = os.path.join(conf.get_lib_dir(), "GoalState.{0}.xml".format(incarnation))
-
-        with open(goal_state_file, "w") as f:
-            f.write(xml_text)
-
-        client = WireClient(WIRESERVER_URL)
-        host = client.get_host_plugin()
-        old_container_id = host.container_id
-
-        # Update the container id
-        new_goal_state = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state.replace("c6d5526c-5ac2-4200-b6e2-56f2b70c5ab2",
-                                                                        "z6d5526c-5ac2-4200-b6e2-56f2b70c5ab2")
-        with patch("azurelinuxagent.common.protocol.wire.WireClient.fetch_config", return_value=new_goal_state):
-            client.update_goal_state(forced=False)
-
-        self.assertNotEqual(old_container_id, host.container_id)
-        self.assertEquals(host.container_id, "z6d5526c-5ac2-4200-b6e2-56f2b70c5ab2")
-
     @patch("azurelinuxagent.common.protocol.wire.WireClient.get_goal_state")
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.get_artifact_request")
     def test_download_ext_handler_pkg_should_not_invoke_host_channel_when_direct_channel_succeeds(self, mock_get_artifact_request, *args):
@@ -735,7 +669,7 @@ class TestWireClient(AgentTestCase):
 
     @patch("azurelinuxagent.common.protocol.wire.WireClient.get_goal_state")
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.get_artifact_request")
-    def test_download_ext_handler_pkg_should_retry_the_host_channel_after_reloading_goal_state(self, mock_get_artifact_request, *args):
+    def test_download_ext_handler_pkg_should_retry_the_host_channel_after_refreshing_host_plugin(self, mock_get_artifact_request, *args):
         mock_get_artifact_request.return_value = "dummy_url", "dummy_header"
         protocol = WireProtocol("foo.bar")
         HostPluginProtocol.set_default_channel(False)
@@ -748,7 +682,7 @@ class TestWireClient(AgentTestCase):
         # As a consequence, goal state 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_goal_state") as mock_update_goal_state:
+            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.stream", wraps=protocol.client.stream) \
                         as patch_direct:
                     with patch("azurelinuxagent.common.protocol.wire.WireProtocol.download_ext_handler_pkg_through_host",
@@ -759,13 +693,13 @@ class TestWireClient(AgentTestCase):
                         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_goal_state.call_count, 1)
+                        self.assertEquals(mock_update_host_plugin_from_goal_state.call_count, 1)
 
                         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_download_ext_handler_pkg_should_update_goal_state_and_not_change_default_channel_if_host_fails(self, mock_get_artifact_request, *args):
+    def test_download_ext_handler_pkg_should_not_change_default_channel_if_host_fails(self, mock_get_artifact_request, *args):
         mock_get_artifact_request.return_value = "dummy_url", "dummy_header"
         protocol = WireProtocol("foo.bar")
         HostPluginProtocol.set_default_channel(False)
@@ -775,7 +709,7 @@ class TestWireClient(AgentTestCase):
 
         # 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_goal_state") as mock_update_goal_state:
+            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.stream", wraps=protocol.client.stream) \
                         as patch_direct:
                     with patch("azurelinuxagent.common.protocol.wire.WireProtocol.download_ext_handler_pkg_through_host",
@@ -786,7 +720,7 @@ class TestWireClient(AgentTestCase):
                         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_goal_state.call_count, 1)
+                        self.assertEquals(mock_update_host_plugin_from_goal_state.call_count, 1)
 
                         self.assertEquals(HostPluginProtocol.is_default_channel(), False)
 
@@ -849,7 +783,7 @@ class TestWireClient(AgentTestCase):
 
     @patch("azurelinuxagent.common.protocol.wire.WireClient.get_goal_state")
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.get_artifact_request")
-    def test_fetch_manifest_should_retry_the_host_channel_after_reloading_goal_state(self, mock_get_artifact_request, *args):
+    def test_fetch_manifest_should_retry_the_host_channel_after_refreshing_the_host_plugin(self, mock_get_artifact_request, *args):
         mock_get_artifact_request.return_value = "dummy_url", "dummy_header"
         client = WireClient("foo.bar")
 
@@ -862,7 +796,7 @@ class TestWireClient(AgentTestCase):
         # As a consequence, goal state 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_goal_state") as mock_update_goal_state:
+            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.fetch_manifest_through_host",
                                wraps=client.fetch_manifest_through_host) as patch_host:
@@ -872,7 +806,7 @@ class TestWireClient(AgentTestCase):
                         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_goal_state.call_count, 1)
+                        self.assertEquals(mock_update_host_plugin_from_goal_state.call_count, 1)
 
                         self.assertEquals(HostPluginProtocol.is_default_channel(), True)
 
@@ -962,7 +896,7 @@ class TestWireClient(AgentTestCase):
 
     @patch("azurelinuxagent.common.protocol.wire.WireClient.get_goal_state")
     @patch("azurelinuxagent.common.protocol.hostplugin.HostPluginProtocol.get_artifact_request")
-    def test_get_artifacts_profile_should_retry_the_host_channel_after_reloading_goal_state(self, mock_get_artifact_request, *args):
+    def test_get_artifacts_profile_should_retry_the_host_channel_after_refreshing_the_host_plugin(self, mock_get_artifact_request, *args):
         mock_get_artifact_request.return_value = "dummy_url", "dummy_header"
         client = WireClient("foo.bar")
         client.ext_conf = ExtensionsConfig(None)
@@ -975,10 +909,10 @@ class TestWireClient(AgentTestCase):
         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, goal state should have been updated and host channel should have been set as default.
+        # 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_goal_state") as mock_update_goal_state:
+            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:
@@ -988,13 +922,13 @@ class TestWireClient(AgentTestCase):
                         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_goal_state.call_count, 1)
+                        self.assertEquals(mock_update_host_plugin_from_goal_state.call_count, 1)
 
                         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_update_goal_state_and_not_change_default_channel_if_host_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, *args):
         mock_get_artifact_request.return_value = "dummy_url", "dummy_header"
         client = WireClient("foo.bar")
         client.ext_conf = ExtensionsConfig(None)
@@ -1007,7 +941,7 @@ class TestWireClient(AgentTestCase):
 
         # 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_goal_state") as mock_update_goal_state:
+            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:
@@ -1017,7 +951,7 @@ class TestWireClient(AgentTestCase):
                         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_goal_state.call_count, 1)
+                        self.assertEquals(mock_update_host_plugin_from_goal_state.call_count, 1)
 
                         self.assertEquals(HostPluginProtocol.is_default_channel(), False)
 
@@ -1114,15 +1048,170 @@ class TestWireClient(AgentTestCase):
 
         # 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_goal_state') as mock_update_goal_state:
+        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)
             self.assertEquals(42, ret)
             self.assertEquals(1, direct_func.counter)
             self.assertEquals(2, host_func.counter)
-            self.assertEquals(1, mock_update_goal_state.call_count)
+            self.assertEquals(1, mock_update_host_plugin_from_goal_state.call_count)
             self.assertEquals(True, client.get_host_plugin().is_default_channel())
 
 
+class GoalStateTestCase(AgentTestCase):
+    """
+    Base class for UpdateGoalStateTestCase and UpdateHostPluginFromGoalStateTestCase
+    """
+    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 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()
+
+            sequence_number = self.wire_client.ext_conf.ext_handlers.extHandlers[0].properties.extensions[0].sequenceNumber
+
+            self.assertEqual(self.wire_client.goal_state.incarnation, new_incarnation)
+            self.assertEqual(self.wire_client.hosting_env.deployment_name, new_hosting_env_deployment_name)
+            self.assertEqual(self.wire_client.shared_conf.xml_text, new_shared_conf)
+            self.assertEqual(sequence_number, new_sequence_number)
+            self.assertEqual(len(self.wire_client.certs.cert_list.certificates), 0)
+
+            self.assertEqual(self.wire_client.host_plugin.container_id, new_container_id)
+            self.assertEqual(self.wire_client.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):
+        # 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.goal_state.xml_text
+        shared_conf = self.wire_client.shared_conf.xml_text
+        container_id = self.wire_client.host_plugin.container_id
+        role_config_name = self.wire_client.host_plugin.role_config_name
+
+        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()))
+
+        self.wire_client.update_goal_state()
+
+        self.assertEqual(self.wire_client.goal_state.xml_text, goal_state)
+        self.assertEqual(self.wire_client.shared_conf.xml_text, shared_conf)
+
+        self.assertEqual(self.wire_client.host_plugin.container_id, container_id)
+        self.assertEqual(self.wire_client.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.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()))
+
+        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
+
+        self.wire_client.update_goal_state(forced=True)
+
+        self.assertEqual(self.wire_client.goal_state.incarnation, incarnation)
+        self.assertEqual(self.wire_client.shared_conf.xml_text, new_shared_conf)
+
+        self.assertEqual(self.wire_client.host_plugin.container_id, new_container_id)
+        self.assertEqual(self.wire_client.host_plugin.role_config_name, new_role_config_name)
+
+
+class UpdateHostPluginFromGoalStateTestCase(GoalStateTestCase):
+    """
+    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
+
+            new_container_id = str(uuid.uuid4())
+            new_role_config_name = 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()))
+
+            self.wire_client.update_host_plugin_from_goal_state()
+
+            self.assertEqual(self.wire_client.host_plugin.container_id, new_container_id)
+            self.assertEqual(self.wire_client.host_plugin.role_config_name, new_role_config_name)
+
+            # it should not update the goal state
+            self.assertEqual(self.wire_client.goal_state.xml_text, goal_state)
+            self.assertEqual(self.wire_client.shared_conf.xml_text, shared_conf)
+
+
 class MockResponse:
     def __init__(self, body, status_code):
         self.body = body