Repository: Azure/WALinuxAgent
Pull request #2713 (2022-12-08 02:53:43)

Log collector should not fetch full goal state with extensions
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description
The Log Collector should not fetch the full goal state with extensions and certificates. This can cause a race condition between the log collector and extension handler when downloading certificates.

This PR adds a GoalStateProperties enum which defines the properties that can be fetched in the goal state and updates fetch_full_wire_server_goal_state to take a list of properties from the goal state to populate. This supports fetching different combinations of propreties in the goal state.

Issue # <!-- if any -->
<!--
Please add an informative description that covers that changes made by the pull request. 
This checklist is used to make sure that common issues in a pull request are addressed.
This will expedite the process of getting your pull request merged and avoid extra work on your part to fix issues discovered during the review process.
-->

---

### PR information
- [ ] The title of the PR is clear and informative.
- [ ] There are a small number of commits, each of which has an informative message. This means that previously merged commits do not appear in the history of the PR. For information on cleaning up the commits in your pull request, [see this page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md).
- [ ] If applicable, the PR references the bug/issue that it fixes in the description.
- [ ] New Unit tests were added for the changes made

### Quality of Code and Contribution Guidelines
- [ ] I have read the [contribution guidelines](https://github.com/Azure/WALinuxAgent/blob/master/.github/CONTRIBUTING.md).

codecov[bot]: # [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2713?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) Report
> Merging [#2713](https://codecov.io/gh/Azure/WALinuxAgent/pull/2713?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (c642900) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/cdb25bfdb4287192d4d19ad559a548cb01f7eff1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (cdb25bf) will **increase** coverage by `0.01%`.
> The diff coverage is `85.07%`.

```diff
@@             Coverage Diff             @@
##           develop    #2713      +/-   ##
===========================================
+ Coverage    71.95%   71.96%   +0.01%     
===========================================
  Files          104      104              
  Lines        15765    15786      +21     
  Branches      2244     2248       +4     
===========================================
+ Hits         11343    11361      +18     
- Misses        3906     3910       +4     
+ Partials       516      515       -1     
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/2713?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/logcollector.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2713/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9sb2djb2xsZWN0b3IucHk=) | `88.00% <25.00%> (-0.31%)` | :arrow_down: |
| [azurelinuxagent/common/protocol/goal\_state.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2713/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC9nb2FsX3N0YXRlLnB5) | `94.06% <86.53%> (-0.70%)` | :arrow_down: |
| [azurelinuxagent/common/protocol/util.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2713/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC91dGlsLnB5) | `80.43% <100.00%> (+0.10%)` | :arrow_up: |
| [azurelinuxagent/common/protocol/wire.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2713/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC93aXJlLnB5) | `78.24% <100.00%> (ø)` | |
| [azurelinuxagent/common/cgroupconfigurator.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2713/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBjb25maWd1cmF0b3IucHk=) | `72.24% <0.00%> (+0.31%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure)


Patch:
diff --git a/azurelinuxagent/common/logcollector.py b/azurelinuxagent/common/logcollector.py
index b0da848f..393333c9 100644
--- a/azurelinuxagent/common/logcollector.py
+++ b/azurelinuxagent/common/logcollector.py
@@ -34,6 +34,7 @@ from azurelinuxagent.common.logcollector_manifests import MANIFEST_NORMAL, MANIF
 
 # Please note: be careful when adding agent dependencies in this module.
 # This module uses its own logger and logs to its own file, not to the agent log.
+from azurelinuxagent.common.protocol.goal_state import GoalStateProperties
 from azurelinuxagent.common.protocol.util import get_protocol_util
 
 _EXTENSION_LOG_DIR = get_ext_log_dir()
@@ -117,8 +118,8 @@ class LogCollector(object):
 
     @staticmethod
     def _initialize_telemetry():
-        protocol = get_protocol_util().get_protocol()
-        protocol.client.update_goal_state(force_update=True)
+        protocol = get_protocol_util().get_protocol(init_goal_state=False)
+        protocol.client.reset_goal_state(goalstate_properties=GoalStateProperties.RoleConfig | GoalStateProperties.HostingEnv)
         # Initialize the common parameters for telemetry events
         initialize_event_logger_vminfo_common_parameters(protocol)
 
diff --git a/azurelinuxagent/common/protocol/goal_state.py b/azurelinuxagent/common/protocol/goal_state.py
index 664b70ef..ed96159c 100644
--- a/azurelinuxagent/common/protocol/goal_state.py
+++ b/azurelinuxagent/common/protocol/goal_state.py
@@ -48,6 +48,18 @@ TRANSPORT_PRV_FILE_NAME = "TransportPrivate.pem"
 _GET_GOAL_STATE_MAX_ATTEMPTS = 6
 
 
+class GoalStateProperties(object):
+    """
+    Enum for defining the properties that we fetch in the goal state
+    """
+    RoleConfig = 0x1
+    HostingEnv = 0x2
+    SharedConfig = 0x4
+    ExtensionsGoalState = 0x8
+    RemoteAccessInfo = 0x10
+    All = RoleConfig | HostingEnv | SharedConfig | ExtensionsGoalState | RemoteAccessInfo
+
+
 class GoalStateInconsistentError(ProtocolError):
     """
     Indicates an inconsistency in the goal state (e.g. missing tenant certificate)
@@ -57,7 +69,7 @@ class GoalStateInconsistentError(ProtocolError):
 
 
 class GoalState(object):
-    def __init__(self, wire_client, silent=False):
+    def __init__(self, wire_client, goal_state_properties=GoalStateProperties.All, silent=False):
         """
         Fetches the goal state using the given wire client.
 
@@ -72,6 +84,7 @@ class GoalState(object):
             self._wire_client = wire_client
             self._history = None
             self._extensions_goal_state = None  # populated from vmSettings or extensionsConfig
+            self._goal_state_properties = goal_state_properties
             self.logger = logger.Logger(logger.DEFAULT_LOGGER)
             self.logger.silent = silent
 
@@ -99,35 +112,59 @@ class GoalState(object):
 
     @property
     def container_id(self):
-        return self._container_id
+        if not self._goal_state_properties & GoalStateProperties.RoleConfig:
+            raise ProtocolError("ContainerId is not in goal state properties")
+        else:
+            return self._container_id
 
     @property
     def role_instance_id(self):
-        return self._role_instance_id
+        if not self._goal_state_properties & GoalStateProperties.RoleConfig:
+            raise ProtocolError("RoleInstanceId is not in goal state properties")
+        else:
+            return self._role_instance_id
 
     @property
     def role_config_name(self):
-        return self._role_config_name
+        if not self._goal_state_properties & GoalStateProperties.RoleConfig:
+            raise ProtocolError("RoleConfig is not in goal state properties")
+        else:
+            return self._role_config_name
 
     @property
     def extensions_goal_state(self):
-        return self._extensions_goal_state
+        if not self._goal_state_properties & GoalStateProperties.ExtensionsGoalState:
+            raise ProtocolError("ExtensionsGoalState is not in goal state properties")
+        else:
+            return self._extensions_goal_state
 
     @property
     def certs(self):
-        return self._certs
+        if not self._goal_state_properties & GoalStateProperties.ExtensionsGoalState:
+            raise ProtocolError("Certificates is not in goal state properties")
+        else:
+            return self._certs
 
     @property
     def hosting_env(self):
-        return self._hosting_env
+        if not self._goal_state_properties & GoalStateProperties.HostingEnv:
+            raise ProtocolError("HostingEnvironment is not in goal state properties")
+        else:
+            return self._hosting_env
 
     @property
     def shared_conf(self):
-        return self._shared_conf
+        if not self._goal_state_properties & GoalStateProperties.SharedConfig:
+            raise ProtocolError("SharedConfig is not in goal state properties")
+        else:
+            return self._shared_conf
 
     @property
     def remote_access(self):
-        return self._remote_access
+        if not self._goal_state_properties & GoalStateProperties.RemoteAccessInfo:
+            raise ProtocolError("RemoteAccessInfo is not in goal state properties")
+        else:
+            return self._remote_access
 
     def fetch_agent_manifest(self, family_name, uris):
         """
@@ -190,11 +227,12 @@ class GoalState(object):
             add_event(op=WALAEventOperation.GoalState, message=message)
 
         vm_settings, vm_settings_updated = None, False
-        try:
-            vm_settings, vm_settings_updated = GoalState._fetch_vm_settings(self._wire_client, force_update=force_update)
-        except VmSettingsSupportStopped as exception:  # If the HGAP stopped supporting vmSettings, we need to use the goal state from the WireServer
-            self._restore_wire_server_goal_state(incarnation, xml_text, xml_doc, exception)
-            return
+        if self._goal_state_properties & GoalStateProperties.ExtensionsGoalState:
+            try:
+                vm_settings, vm_settings_updated = GoalState._fetch_vm_settings(self._wire_client, force_update=force_update)
+            except VmSettingsSupportStopped as exception:  # If the HGAP stopped supporting vmSettings, we need to use the goal state from the WireServer
+                self._restore_wire_server_goal_state(incarnation, xml_text, xml_doc, exception)
+                return
 
         if vm_settings_updated:
             self.logger.info('')
@@ -356,40 +394,48 @@ class GoalState(object):
             self.logger.info(message)
             add_event(op=WALAEventOperation.GoalState, message=message)
 
-            role_instance = find(xml_doc, "RoleInstance")
-            role_instance_id = findtext(role_instance, "InstanceId")
-            role_config = find(role_instance, "Configuration")
-            role_config_name = findtext(role_config, "ConfigName")
-            container = find(xml_doc, "Container")
-            container_id = findtext(container, "ContainerId")
+            role_instance_id = None
+            role_config_name = None
+            container_id = None
+            if GoalStateProperties.RoleConfig & self._goal_state_properties:
+                role_instance = find(xml_doc, "RoleInstance")
+                role_instance_id = findtext(role_instance, "InstanceId")
+                role_config = find(role_instance, "Configuration")
+                role_config_name = findtext(role_config, "ConfigName")
+                container = find(xml_doc, "Container")
+                container_id = findtext(container, "ContainerId")
 
             extensions_config_uri = findtext(xml_doc, "ExtensionsConfig")
-            if extensions_config_uri is None:
+            if not (GoalStateProperties.ExtensionsGoalState & self._goal_state_properties) or extensions_config_uri is None:
                 extensions_config = ExtensionsGoalStateFactory.create_empty(incarnation)
             else:
                 xml_text = self._wire_client.fetch_config(extensions_config_uri, self._wire_client.get_header())
                 extensions_config = ExtensionsGoalStateFactory.create_from_extensions_config(incarnation, xml_text, self._wire_client)
                 self._history.save_extensions_config(extensions_config.get_redacted_text())
 
-            hosting_env_uri = findtext(xml_doc, "HostingEnvironmentConfig")
-            xml_text = self._wire_client.fetch_config(hosting_env_uri, self._wire_client.get_header())
-            hosting_env = HostingEnv(xml_text)
-            self._history.save_hosting_env(xml_text)
-
-            shared_conf_uri = findtext(xml_doc, "SharedConfig")
-            xml_text = self._wire_client.fetch_config(shared_conf_uri, self._wire_client.get_header())
-            shared_config = SharedConfig(xml_text)
-            self._history.save_shared_conf(xml_text)
-            # SharedConfig.xml is used by other components (Azsec and Singularity/HPC Infiniband), so save it to the agent's root directory as well
-            shared_config_file = os.path.join(conf.get_lib_dir(), SHARED_CONF_FILE_NAME)
-            try:
-                fileutil.write_file(shared_config_file, xml_text)
-            except Exception as e:
-                logger.warn("Failed to save {0}: {1}".format(shared_config, e))
+            hosting_env = None
+            if GoalStateProperties.HostingEnv & self._goal_state_properties:
+                hosting_env_uri = findtext(xml_doc, "HostingEnvironmentConfig")
+                xml_text = self._wire_client.fetch_config(hosting_env_uri, self._wire_client.get_header())
+                hosting_env = HostingEnv(xml_text)
+                self._history.save_hosting_env(xml_text)
+
+            shared_config = None
+            if GoalStateProperties.SharedConfig & self._goal_state_properties:
+                shared_conf_uri = findtext(xml_doc, "SharedConfig")
+                xml_text = self._wire_client.fetch_config(shared_conf_uri, self._wire_client.get_header())
+                shared_config = SharedConfig(xml_text)
+                self._history.save_shared_conf(xml_text)
+                # SharedConfig.xml is used by other components (Azsec and Singularity/HPC Infiniband), so save it to the agent's root directory as well
+                shared_config_file = os.path.join(conf.get_lib_dir(), SHARED_CONF_FILE_NAME)
+                try:
+                    fileutil.write_file(shared_config_file, xml_text)
+                except Exception as e:
+                    logger.warn("Failed to save {0}: {1}".format(shared_config, e))
 
             certs = EmptyCertificates()
             certs_uri = findtext(xml_doc, "Certificates")
-            if certs_uri is not None:
+            if (GoalStateProperties.ExtensionsGoalState & self._goal_state_properties) and certs_uri is not None:
                 xml_text = self._wire_client.fetch_config(certs_uri, self._wire_client.get_header_for_cert())
                 certs = Certificates(xml_text, self.logger)
                 # Log and save the certificates summary (i.e. the thumbprint but not the certificate itself) to the goal state history
@@ -403,11 +449,12 @@ class GoalState(object):
                 self._history.save_certificates(json.dumps(certs.summary))
 
             remote_access = None
-            remote_access_uri = findtext(container, "RemoteAccessInfo")
-            if remote_access_uri is not None:
-                xml_text = self._wire_client.fetch_config(remote_access_uri, self._wire_client.get_header_for_cert())
-                remote_access = RemoteAccess(xml_text)
-                self._history.save_remote_access(xml_text)
+            if GoalStateProperties.RemoteAccessInfo & self._goal_state_properties:
+                remote_access_uri = findtext(container, "RemoteAccessInfo")
+                if remote_access_uri is not None:
+                    xml_text = self._wire_client.fetch_config(remote_access_uri, self._wire_client.get_header_for_cert())
+                    remote_access = RemoteAccess(xml_text)
+                    self._history.save_remote_access(xml_text)
 
             self._incarnation = incarnation
             self._role_instance_id = role_instance_id
diff --git a/azurelinuxagent/common/protocol/util.py b/azurelinuxagent/common/protocol/util.py
index 92b691e9..7d7f9016 100644
--- a/azurelinuxagent/common/protocol/util.py
+++ b/azurelinuxagent/common/protocol/util.py
@@ -188,7 +188,7 @@ class ProtocolUtil(SingletonPerThread):
                 return
             logger.error("Failed to clear wiresever endpoint: {0}", e)
 
-    def _detect_protocol(self):
+    def _detect_protocol(self, init_goal_state=True):
         """
         Probe protocol endpoints in turn.
         """
@@ -217,7 +217,7 @@ class ProtocolUtil(SingletonPerThread):
 
                 try:
                     protocol = WireProtocol(endpoint)
-                    protocol.detect()
+                    protocol.detect(init_goal_state=init_goal_state)
                     self._set_wireserver_endpoint(endpoint)
                     return protocol
 
@@ -268,7 +268,7 @@ class ProtocolUtil(SingletonPerThread):
         finally:
             self._lock.release()
 
-    def get_protocol(self):
+    def get_protocol(self, init_goal_state=True):
         """
         Detect protocol by endpoint.
         :returns: protocol instance
@@ -296,7 +296,7 @@ class ProtocolUtil(SingletonPerThread):
 
             logger.info("Detect protocol endpoint")
 
-            protocol = self._detect_protocol()
+            protocol = self._detect_protocol(init_goal_state=init_goal_state)
 
             IOErrorCounter.set_protocol_endpoint(endpoint=protocol.get_endpoint())
             self._save_protocol(WIRE_PROTOCOL_NAME)
diff --git a/azurelinuxagent/common/protocol/wire.py b/azurelinuxagent/common/protocol/wire.py
index d5aeda71..38a3e062 100644
--- a/azurelinuxagent/common/protocol/wire.py
+++ b/azurelinuxagent/common/protocol/wire.py
@@ -37,7 +37,8 @@ from azurelinuxagent.common.event import add_event, WALAEventOperation, report_e
 from azurelinuxagent.common.exception import ProtocolNotFoundError, \
     ResourceGoneError, ExtensionDownloadError, InvalidContainerError, ProtocolError, HttpError, ExtensionErrorCodes
 from azurelinuxagent.common.future import httpclient, bytebuffer, ustr
-from azurelinuxagent.common.protocol.goal_state import GoalState, TRANSPORT_CERT_FILE_NAME, TRANSPORT_PRV_FILE_NAME
+from azurelinuxagent.common.protocol.goal_state import GoalState, TRANSPORT_CERT_FILE_NAME, TRANSPORT_PRV_FILE_NAME, \
+    GoalStateProperties
 from azurelinuxagent.common.protocol.hostplugin import HostPluginProtocol
 from azurelinuxagent.common.protocol.restapi import DataContract, ProvisionStatus, VMInfo, VMStatus
 from azurelinuxagent.common.telemetryevent import GuestAgentExtensionEventsSchema
@@ -72,7 +73,7 @@ class WireProtocol(DataContract):
             raise ProtocolError("WireProtocol endpoint is None")
         self.client = WireClient(endpoint)
 
-    def detect(self):
+    def detect(self, init_goal_state=True):
         self.client.check_wire_protocol_version()
 
         trans_prv_file = os.path.join(conf.get_lib_dir(),
@@ -83,11 +84,9 @@ class WireProtocol(DataContract):
         cryptutil.gen_transport_cert(trans_prv_file, trans_cert_file)
 
         # Initialize the goal state, including all the inner properties
-        logger.info('Initializing goal state during protocol detection')
-        self.client.update_goal_state(force_update=True)
-
-    def update_goal_state(self, silent=False):
-        self.client.update_goal_state(silent=silent)
+        if init_goal_state:
+            logger.info('Initializing goal state during protocol detection')
+            self.client.reset_goal_state()
 
     def update_host_plugin_from_goal_state(self):
         self.client.update_host_plugin_from_goal_state()
@@ -778,15 +777,12 @@ class WireClient(object):
             self._host_plugin.update_container_id(container_id)
             self._host_plugin.update_role_config_name(role_config_name)
 
-    def update_goal_state(self, force_update=False, silent=False):
+    def update_goal_state(self, silent=False):
         """
-        Updates the goal state if the incarnation or etag changed or if 'force_update' is True
+        Updates the goal state if the incarnation or etag changed
         """
         try:
-            if force_update and not silent:
-                logger.info("Forcing an update of the goal state.")
-
-            if self._goal_state is None or force_update:
+            if self._goal_state is None:
                 self._goal_state = GoalState(self, silent=silent)
             else:
                 self._goal_state.update(silent=silent)
@@ -796,6 +792,21 @@ class WireClient(object):
         except Exception as exception:
             raise ProtocolError("Error fetching goal state: {0}".format(ustr(exception)))
 
+    def reset_goal_state(self, goal_state_properties=GoalStateProperties.All, silent=False):
+        """
+        Resets the goal state
+        """
+        try:
+            if not silent:
+                logger.info("Forcing an update of the goal state.")
+
+            self._goal_state = GoalState(self, goal_state_properties=goal_state_properties, silent=silent)
+
+        except ProtocolError:
+            raise
+        except Exception as exception:
+            raise ProtocolError("Error fetching goal state: {0}".format(ustr(exception)))
+
     def get_goal_state(self):
         if self._goal_state is None:
             raise ProtocolError("Trying to fetch goal state before initialization!")
@@ -925,7 +936,7 @@ class WireClient(object):
 
         if extensions_goal_state.status_upload_blob is None:
             # the status upload blob is in ExtensionsConfig so force a full goal state refresh
-            self.update_goal_state(force_update=True, silent=True)
+            self.reset_goal_state(silent=True)
             extensions_goal_state = self.get_goal_state().extensions_goal_state
 
             if extensions_goal_state.status_upload_blob is None:
diff --git a/azurelinuxagent/daemon/main.py b/azurelinuxagent/daemon/main.py
index c608768a..1eb58ec9 100644
--- a/azurelinuxagent/daemon/main.py
+++ b/azurelinuxagent/daemon/main.py
@@ -28,6 +28,7 @@ import azurelinuxagent.common.utils.fileutil as fileutil
 from azurelinuxagent.common.event import add_event, WALAEventOperation, initialize_event_logger_vminfo_common_parameters
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.osutil import get_osutil
+from azurelinuxagent.common.protocol.goal_state import GoalState, GoalStateProperties
 from azurelinuxagent.common.protocol.util import get_protocol_util
 from azurelinuxagent.common.rdma import setup_rdma_device
 from azurelinuxagent.common.utils import textutil
@@ -160,9 +161,9 @@ class DaemonHandler(object):
                 #   current values.
                 protocol = self.protocol_util.get_protocol()
 
-                protocol.client.update_goal_state(force_update=True)
+                goal_state = GoalState(protocol, goal_state_properties=GoalStateProperties.SharedConfig)
 
-                setup_rdma_device(nd_version, protocol.client.get_shared_conf())
+                setup_rdma_device(nd_version, goal_state.shared_conf)
             except Exception as e:
                 logger.error("Error setting up rdma device: %s" % e)
         else:
diff --git a/azurelinuxagent/ga/update.py b/azurelinuxagent/ga/update.py
index 8d2c97df..cd758b97 100644
--- a/azurelinuxagent/ga/update.py
+++ b/azurelinuxagent/ga/update.py
@@ -485,7 +485,7 @@ class UpdateHandler(object):
         try:
             max_errors_to_log = 3
 
-            protocol.update_goal_state(silent=self._update_goal_state_error_count >= max_errors_to_log)
+            protocol.client.update_goal_state(silent=self._update_goal_state_error_count >= max_errors_to_log)
 
             self._goal_state = protocol.get_goal_state()
 


Test patch:
diff --git a/tests/common/test_event.py b/tests/common/test_event.py
index 7191f3c3..fad21155 100644
--- a/tests/common/test_event.py
+++ b/tests/common/test_event.py
@@ -161,12 +161,12 @@ class TestEvent(HttpRequestPredicates, AgentTestCase):
             self.assertEqual(contained_id, 'c6d5526c-5ac2-4200-b6e2-56f2b70c5ab2', "Incorrect container ID")
 
             protocol.mock_wire_data.set_container_id('AAAAAAAA-BBBB-CCCC-DDDD-EEEEEEEEEEEE')
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
             contained_id = create_event_and_return_container_id()
             self.assertEqual(contained_id, 'AAAAAAAA-BBBB-CCCC-DDDD-EEEEEEEEEEEE', "Incorrect container ID")
 
             protocol.mock_wire_data.set_container_id('11111111-2222-3333-4444-555555555555')
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
             contained_id = create_event_and_return_container_id()
             self.assertEqual(contained_id, '11111111-2222-3333-4444-555555555555', "Incorrect container ID")
 
diff --git a/tests/ga/test_extension.py b/tests/ga/test_extension.py
index 5947e9e6..2272a190 100644
--- a/tests/ga/test_extension.py
+++ b/tests/ga/test_extension.py
@@ -510,7 +510,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         test_data.set_incarnation(2)
         test_data.set_extensions_config_version("1.0.1")
         test_data.set_manifest_version('1.0.1')
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         # Ensure the patched command fails
         patch_command.return_value = "exit 1"
@@ -542,7 +542,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Test goal state changed
         test_data.set_incarnation(2)
         test_data.set_extensions_config_sequence_number(1)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -555,7 +555,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         test_data.set_incarnation(3)
         test_data.set_extensions_config_version("1.1.1")
         test_data.set_extensions_config_sequence_number(2)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -567,7 +567,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         test_data.set_incarnation(4)
         test_data.set_extensions_config_version("1.2.0")
         test_data.set_extensions_config_sequence_number(3)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -578,7 +578,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Test disable
         test_data.set_incarnation(5)
         test_data.set_extensions_config_state(ExtensionRequestedState.Disabled)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -588,7 +588,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Test uninstall
         test_data.set_incarnation(6)
         test_data.set_extensions_config_state(ExtensionRequestedState.Uninstall)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -597,7 +597,7 @@ class TestExtension_Deprecated(TestExtensionBase):
 
         # Test uninstall again!
         test_data.set_incarnation(7)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -620,7 +620,7 @@ class TestExtension_Deprecated(TestExtensionBase):
 
         # Update Incarnation
         test_data.set_incarnation(2)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -720,7 +720,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Uninstall the Plugin and make sure Disable called
         test_data.set_incarnation(2)
         test_data.set_extensions_config_state(ExtensionRequestedState.Uninstall)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         with enable_invocations(test_ext) as invocation_record:
             exthandlers_handler.run()
@@ -799,7 +799,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         dep_ext_level_4 = extension_emulator(name="OSTCExtensions.OtherExampleHandlerLinux")
         test_data.ext_conf = test_data.ext_conf.replace("dependencyLevel=\"2\"", "dependencyLevel=\"3\"")
         test_data.ext_conf = test_data.ext_conf.replace("dependencyLevel=\"1\"", "dependencyLevel=\"4\"")
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         with enable_invocations(dep_ext_level_3, dep_ext_level_4) as invocation_record:
             exthandlers_handler.run()
@@ -826,7 +826,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # the last one disabled.
         test_data.set_incarnation(3)
         test_data.set_extensions_config_state(ExtensionRequestedState.Disabled)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         with enable_invocations(dep_ext_level_3, dep_ext_level_4) as invocation_record:
             exthandlers_handler.run()
@@ -858,7 +858,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         dep_ext_level_6 = extension_emulator(name="OSTCExtensions.ExampleHandlerLinux")
         test_data.ext_conf = test_data.ext_conf.replace("dependencyLevel=\"3\"", "dependencyLevel=\"6\"")
         test_data.ext_conf = test_data.ext_conf.replace("dependencyLevel=\"4\"", "dependencyLevel=\"5\"")
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         with enable_invocations(dep_ext_level_5, dep_ext_level_6) as invocation_record:
             exthandlers_handler.run()
@@ -938,7 +938,7 @@ class TestExtension_Deprecated(TestExtensionBase):
                 _assert_event_reported_only_on_incarnation_change(expected_count=1)
 
                 test_data.set_incarnation(2)
-                protocol.update_goal_state()
+                protocol.client.update_goal_state()
 
                 exthandlers_handler.run()
                 exthandlers_handler.report_ext_handlers_status()
@@ -949,7 +949,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Test it recovers on a new goal state if Handler succeeds
         test_data.set_incarnation(3)
         test_data.set_extensions_config_sequence_number(1)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -961,7 +961,7 @@ class TestExtension_Deprecated(TestExtensionBase):
 
         # Update incarnation to confirm extension invocation order
         test_data.set_incarnation(4)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         dep_ext_level_2 = extension_emulator(name="OSTCExtensions.ExampleHandlerLinux")
         dep_ext_level_1 = extension_emulator(name="OSTCExtensions.OtherExampleHandlerLinux")
@@ -1023,7 +1023,7 @@ class TestExtension_Deprecated(TestExtensionBase):
 
         # Test goal state changed
         test_data.set_incarnation(2)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -1034,7 +1034,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Test minor version bump
         test_data.set_incarnation(3)
         test_data.set_extensions_config_version("1.1.0")
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -1045,7 +1045,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Test hotfix version bump
         test_data.set_incarnation(4)
         test_data.set_extensions_config_version("1.1.1")
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -1056,7 +1056,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Test disable
         test_data.set_incarnation(5)
         test_data.set_extensions_config_state(ExtensionRequestedState.Disabled)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -1066,7 +1066,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Test uninstall
         test_data.set_incarnation(6)
         test_data.set_extensions_config_state(ExtensionRequestedState.Uninstall)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -1075,7 +1075,7 @@ class TestExtension_Deprecated(TestExtensionBase):
 
         # Test uninstall again!
         test_data.set_incarnation(7)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -1085,7 +1085,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Test re-install
         test_data.set_incarnation(8)
         test_data.set_extensions_config_state(ExtensionRequestedState.Enabled)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -1096,7 +1096,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Test version bump post-re-install
         test_data.set_incarnation(9)
         test_data.set_extensions_config_version("1.2.0")
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -1107,7 +1107,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Test rollback
         test_data.set_incarnation(10)
         test_data.set_extensions_config_version("1.1.0")
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -1170,7 +1170,7 @@ class TestExtension_Deprecated(TestExtensionBase):
             # Uninstall extensions now
             test_data.set_extensions_config_state(ExtensionRequestedState.Uninstall)
             test_data.set_incarnation(2)
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
 
             exthandlers_handler.run()
             exthandlers_handler.report_ext_handlers_status()
@@ -1192,7 +1192,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Since the installed version is not in PIR anymore, we need to also remove it from manifest file
         test_data.manifest = test_data.manifest.replace("1.0.0", "9.9.9")
         test_data.set_incarnation(2)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -1666,7 +1666,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         test_data.set_incarnation(2)
         test_data.set_extensions_config_version("1.0.1")
         test_data.set_manifest_version('1.0.1')
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         # Ensure new extension can be enabled
         exthandlers_handler.run()
@@ -1753,7 +1753,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Next incarnation, disable extension
         test_data.set_incarnation(2)
         test_data.set_extensions_config_state(ExtensionRequestedState.Disabled)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -1784,7 +1784,7 @@ class TestExtension_Deprecated(TestExtensionBase):
         # Next incarnation, disable extension
         test_data.set_incarnation(2)
         test_data.set_extensions_config_state(ExtensionRequestedState.Uninstall)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
 
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
@@ -1841,7 +1841,7 @@ class TestExtension_Deprecated(TestExtensionBase):
             # If the incarnation number changes (there's a new goal state), ensure we go through the entire upgrade
             # process again.
             test_data.set_incarnation(3)
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
 
             exthandlers_handler.run()
             exthandlers_handler.report_ext_handlers_status()
@@ -1888,7 +1888,7 @@ class TestExtension_Deprecated(TestExtensionBase):
 
             # Force a new goal state incarnation, only then will we attempt the upgrade again
             test_data.set_incarnation(3)
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
 
             # Ensure disable won't fail by making launch_command a no-op
             with patch('azurelinuxagent.ga.exthandlers.ExtHandlerInstance.launch_command') as patch_launch_command:  # pylint: disable=unused-variable
@@ -2111,7 +2111,7 @@ class TestExtension_Deprecated(TestExtensionBase):
             # Initiating another run which shouldn't have any failed env variables in it if no failures
             # Updating Incarnation
             test_data.set_incarnation(3)
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
 
             exthandlers_handler.run()
             exthandlers_handler.report_ext_handlers_status()
@@ -3355,7 +3355,7 @@ class TestExtension(TestExtensionBase, HttpRequestPredicates):
 
                 # Update GoalState
                 protocol.mock_wire_data.set_incarnation(2)
-                protocol.update_goal_state()
+                protocol.client.update_goal_state()
 
                 with patch.object(conf, 'get_extensions_enabled', return_value=False):
                     assert_extensions_called(exthandlers_handler, expected_call_count=0)
@@ -3369,7 +3369,7 @@ class TestExtension(TestExtensionBase, HttpRequestPredicates):
 
                 # Enabled on_hold property in artifact_blob
                 mock_in_vm_artifacts_profile_response = MockHttpResponse(200, body='{ "onHold": true }'.encode('utf-8'))
-                protocol.client.update_goal_state(force_update=True)
+                protocol.client.reset_goal_state()
 
                 with patch.object(conf, 'get_extensions_enabled', return_value=True):
                     with patch.object(conf, "get_enable_overprovisioning", return_value=True):
@@ -3377,7 +3377,7 @@ class TestExtension(TestExtensionBase, HttpRequestPredicates):
 
                 # Disabled on_hold property in artifact_blob
                 mock_in_vm_artifacts_profile_response = MockHttpResponse(200, body='{ "onHold": false }'.encode('utf-8'))
-                protocol.client.update_goal_state(force_update=True)
+                protocol.client.reset_goal_state()
 
                 with patch.object(conf, 'get_extensions_enabled', return_value=True):
                     with patch.object(conf, "get_enable_overprovisioning", return_value=True):
@@ -3410,7 +3410,7 @@ class TestExtension(TestExtensionBase, HttpRequestPredicates):
                         return None
                     protocol.set_http_handlers(http_get_handler=http_get_handler)
 
-                    protocol.client.update_goal_state(force_update=True)
+                    protocol.client.reset_goal_state()
 
                     exthandlers_handler.run()
                     exthandlers_handler.report_ext_handlers_status()
diff --git a/tests/ga/test_monitor.py b/tests/ga/test_monitor.py
index d5700bc9..5853b23e 100644
--- a/tests/ga/test_monitor.py
+++ b/tests/ga/test_monitor.py
@@ -188,7 +188,7 @@ class TestExtensionMetricsDataTelemetry(AgentTestCase):
         CGroupsTelemetry.reset()
         clear_singleton_instances(ProtocolUtil)
         protocol = WireProtocol('endpoint')
-        protocol.update_goal_state = MagicMock()
+        protocol.client.update_goal_state = MagicMock()
         self.get_protocol = patch('azurelinuxagent.common.protocol.util.ProtocolUtil.get_protocol', return_value=protocol)
         self.get_protocol.start()
 
diff --git a/tests/ga/test_multi_config_extension.py b/tests/ga/test_multi_config_extension.py
index a9a07bd6..365052f5 100644
--- a/tests/ga/test_multi_config_extension.py
+++ b/tests/ga/test_multi_config_extension.py
@@ -253,7 +253,7 @@ class TestMultiConfigExtensions(_MultiConfigBaseTestClass):
         self.test_data['ext_conf'] = os.path.join(self._MULTI_CONFIG_TEST_DATA, 'ext_conf_mc_disabled_extensions.xml')
         protocol.mock_wire_data = WireProtocolData(self.test_data)
         protocol.mock_wire_data.set_incarnation(2)
-        protocol.update_goal_state()
+        protocol.client.update_goal_state()
         exthandlers_handler.run()
         exthandlers_handler.report_ext_handlers_status()
 
@@ -319,7 +319,7 @@ class TestMultiConfigExtensions(_MultiConfigBaseTestClass):
             # Case 3: Uninstall Multi-config handler (with enabled extensions) and single config extension
             protocol.mock_wire_data.set_incarnation(3)
             protocol.mock_wire_data.set_extensions_config_state(ExtensionRequestedState.Uninstall)
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
             exthandlers_handler.run()
             exthandlers_handler.report_ext_handlers_status()
             self.assertEqual(0, len(protocol.aggregate_status['aggregateStatus']['handlerAggregateStatus']),
@@ -333,7 +333,7 @@ class TestMultiConfigExtensions(_MultiConfigBaseTestClass):
             failing_version = "19.12.1221"
             protocol.mock_wire_data.set_extensions_config_version(failing_version)
             protocol.mock_wire_data.set_incarnation(2)
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
             exthandlers_handler.run()
             exthandlers_handler.report_ext_handlers_status()
             self.assertEqual(no_of_extensions,
@@ -411,7 +411,7 @@ class TestMultiConfigExtensions(_MultiConfigBaseTestClass):
             self.test_data['ext_conf'] = os.path.join(self._MULTI_CONFIG_TEST_DATA, 'ext_conf_mc_update_extensions.xml')
             protocol.mock_wire_data = WireProtocolData(self.test_data)
             protocol.mock_wire_data.set_incarnation(2)
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
 
             new_version = "1.1.0"
             new_first_ext = extension_emulator(name="OSTCExtensions.ExampleHandlerLinux.firstExtension",
@@ -460,7 +460,7 @@ class TestMultiConfigExtensions(_MultiConfigBaseTestClass):
             self.test_data['ext_conf'] = os.path.join(self._MULTI_CONFIG_TEST_DATA, 'ext_conf_mc_update_extensions.xml')
             protocol.mock_wire_data = WireProtocolData(self.test_data)
             protocol.mock_wire_data.set_incarnation(2)
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
 
             new_version = "1.1.0"
             _, fail_action = Actions.generate_unique_fail()
@@ -529,7 +529,7 @@ class TestMultiConfigExtensions(_MultiConfigBaseTestClass):
             self.test_data['ext_conf'] = os.path.join(self._MULTI_CONFIG_TEST_DATA, 'ext_conf_mc_update_extensions.xml')
             protocol.mock_wire_data = WireProtocolData(self.test_data)
             protocol.mock_wire_data.set_incarnation(2)
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
 
             new_version = "1.1.0"
             fail_code, fail_action = Actions.generate_unique_fail()
@@ -655,7 +655,7 @@ class TestMultiConfigExtensions(_MultiConfigBaseTestClass):
             self.test_data['ext_conf'] = os.path.join(self._MULTI_CONFIG_TEST_DATA, 'ext_conf_mc_disabled_extensions.xml')
             protocol.mock_wire_data = WireProtocolData(self.test_data)
             protocol.mock_wire_data.set_incarnation(2)
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
 
             ext_handler.run()
             ext_handler.report_ext_handlers_status()
@@ -781,7 +781,7 @@ class TestMultiConfigExtensions(_MultiConfigBaseTestClass):
                                                           'ext_conf_mc_update_extensions.xml')
                 protocol.mock_wire_data = WireProtocolData(self.test_data)
                 protocol.mock_wire_data.set_incarnation(2)
-                protocol.update_goal_state()
+                protocol.client.update_goal_state()
                 exthandlers_handler.run()
                 exthandlers_handler.report_ext_handlers_status()
 
@@ -961,7 +961,7 @@ class TestMultiConfigExtensions(_MultiConfigBaseTestClass):
             self.test_data['ext_conf'] = "wire/ext_conf_required_features.xml"
             protocol.mock_wire_data = WireProtocolData(self.test_data)
             protocol.mock_wire_data.set_incarnation(2)
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
             # Assert the extension status is the same as we reported for Incarnation 1.
             self.__run_and_assert_generic_case(exthandlers_handler, protocol, no_of_extensions=4, with_message=False)
 
@@ -1021,7 +1021,7 @@ class TestMultiConfigExtensions(_MultiConfigBaseTestClass):
         with self.__setup_generic_test_env() as (exthandlers_handler, protocol, old_exts):
 
             protocol.mock_wire_data.set_incarnation(2)
-            protocol.update_goal_state()
+            protocol.client.update_goal_state()
 
             # Mock manifest to not support multiple extensions
             with patch('azurelinuxagent.ga.exthandlers.HandlerManifest.supports_multiple_extensions', return_value=False):
diff --git a/tests/protocol/test_goal_state.py b/tests/protocol/test_goal_state.py
index d853363c..869da68c 100644
--- a/tests/protocol/test_goal_state.py
+++ b/tests/protocol/test_goal_state.py
@@ -14,7 +14,8 @@ from azurelinuxagent.common.protocol.extensions_goal_state import GoalStateSourc
 from azurelinuxagent.common.protocol.extensions_goal_state_from_extensions_config import ExtensionsGoalStateFromExtensionsConfig
 from azurelinuxagent.common.protocol.extensions_goal_state_from_vm_settings import ExtensionsGoalStateFromVmSettings
 from azurelinuxagent.common.protocol import hostplugin
-from azurelinuxagent.common.protocol.goal_state import GoalState, GoalStateInconsistentError, _GET_GOAL_STATE_MAX_ATTEMPTS
+from azurelinuxagent.common.protocol.goal_state import GoalState, GoalStateInconsistentError, \
+    _GET_GOAL_STATE_MAX_ATTEMPTS, GoalStateProperties
 from azurelinuxagent.common.exception import ProtocolError
 from azurelinuxagent.common.utils import fileutil
 from azurelinuxagent.common.utils.archive import ARCHIVE_DIRECTORY_NAME
@@ -419,3 +420,74 @@ class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
                 for settings in extension.settings:
                     if settings.protectedSettings is not None:
                         self.assertIn(settings.certificateThumbprint, thumbprints, "Certificate is missing from the goal state.")
+
+    def test_it_should_raise_when_goal_state_properties_not_initialized(self):
+        with GoalStateTestCase._create_protocol_ws_and_hgap_in_sync() as protocol:
+            goal_state = GoalState(
+                protocol.client,
+                goal_state_properties=~GoalStateProperties.All)
+
+            goal_state.update()
+
+            with self.assertRaises(ProtocolError) as context:
+                _ = goal_state.container_id
+
+            expected_message = "ContainerId is not in goal state properties"
+            self.assertIn(expected_message, str(context.exception))
+
+            with self.assertRaises(ProtocolError) as context:
+                _ = goal_state.role_config_name
+
+            expected_message = "RoleConfig is not in goal state properties"
+            self.assertIn(expected_message, str(context.exception))
+
+            with self.assertRaises(ProtocolError) as context:
+                _ = goal_state.role_instance_id
+
+            expected_message = "RoleInstanceId is not in goal state properties"
+            self.assertIn(expected_message, str(context.exception))
+
+            with self.assertRaises(ProtocolError) as context:
+                _ = goal_state.extensions_goal_state
+
+            expected_message = "ExtensionsGoalState is not in goal state properties"
+            self.assertIn(expected_message, str(context.exception))
+
+            with self.assertRaises(ProtocolError) as context:
+                _ = goal_state.hosting_env
+
+            expected_message = "HostingEnvironment is not in goal state properties"
+            self.assertIn(expected_message, str(context.exception))
+
+            with self.assertRaises(ProtocolError) as context:
+                _ = goal_state.certs
+
+            expected_message = "Certificates is not in goal state properties"
+            self.assertIn(expected_message, str(context.exception))
+
+            with self.assertRaises(ProtocolError) as context:
+                _ = goal_state.shared_conf
+
+            expected_message = "SharedConfig is not in goal state properties"
+            self.assertIn(expected_message, str(context.exception))
+
+            with self.assertRaises(ProtocolError) as context:
+                _ = goal_state.remote_access
+
+            expected_message = "RemoteAccessInfo is not in goal state properties"
+            self.assertIn(expected_message, str(context.exception))
+
+            goal_state = GoalState(
+                protocol.client,
+                goal_state_properties=GoalStateProperties.All & ~GoalStateProperties.HostingEnv)
+
+            goal_state.update()
+
+            _ = goal_state.container_id, goal_state.role_instance_id, goal_state.role_config_name, \
+                goal_state.extensions_goal_state, goal_state.certs, goal_state.shared_conf, goal_state.remote_access
+
+            with self.assertRaises(ProtocolError) as context:
+                _ = goal_state.hosting_env
+
+            expected_message = "HostingEnvironment is not in goal state properties"
+            self.assertIn(expected_message, str(context.exception))
diff --git a/tests/protocol/test_hostplugin.py b/tests/protocol/test_hostplugin.py
index c980433b..b85ed757 100644
--- a/tests/protocol/test_hostplugin.py
+++ b/tests/protocol/test_hostplugin.py
@@ -244,7 +244,7 @@ class TestHostPlugin(HttpRequestPredicates, AgentTestCase):
         with self.create_mock_protocol() as wire_protocol:
             wire.HostPluginProtocol.is_default_channel = False
 
-            wire_protocol.update_goal_state()
+            wire_protocol.client.update_goal_state()
 
             # act
             wire_protocol.client.upload_status_blob()
@@ -277,7 +277,7 @@ class TestHostPlugin(HttpRequestPredicates, AgentTestCase):
         with self.create_mock_protocol() as wire_protocol:
             wire.HostPluginProtocol.is_default_channel = False
 
-            wire_protocol.update_goal_state()
+            wire_protocol.client.update_goal_state()
 
             # act
             wire_protocol.client.upload_status_blob()
@@ -311,7 +311,7 @@ class TestHostPlugin(HttpRequestPredicates, AgentTestCase):
         with self.create_mock_protocol() as wire_protocol:
             wire.HostPluginProtocol.is_default_channel = False
 
-            wire_protocol.update_goal_state()
+            wire_protocol.client.update_goal_state()
 
             # act
             wire_protocol.client.upload_status_blob()
@@ -345,7 +345,7 @@ class TestHostPlugin(HttpRequestPredicates, AgentTestCase):
         with self.create_mock_protocol() as wire_protocol:
             wire.HostPluginProtocol.is_default_channel = False
 
-            wire_protocol.update_goal_state()
+            wire_protocol.client.update_goal_state()
 
             # act
             self.assertRaises(wire.ProtocolError, wire_protocol.client.upload_status_blob)
diff --git a/tests/protocol/test_wire.py b/tests/protocol/test_wire.py
index 7c121f44..bbf018fc 100644
--- a/tests/protocol/test_wire.py
+++ b/tests/protocol/test_wire.py
@@ -30,6 +30,7 @@ from azurelinuxagent.common.event import WALAEventOperation
 from azurelinuxagent.common.exception import ResourceGoneError, ProtocolError, \
     ExtensionDownloadError, HttpError
 from azurelinuxagent.common.protocol.extensions_goal_state_from_extensions_config import ExtensionsGoalStateFromExtensionsConfig
+from azurelinuxagent.common.protocol.goal_state import GoalStateProperties
 from azurelinuxagent.common.protocol.hostplugin import HostPluginProtocol
 from azurelinuxagent.common.protocol.wire import WireProtocol, WireClient, \
     StatusBlob, VMStatus
@@ -271,23 +272,23 @@ class TestWireProtocol(AgentTestCase, HttpRequestPredicates):
             protocol.set_http_handlers(http_get_handler=http_get_handler)
 
             mock_response = MockHttpResponse(200, body=None)
-            protocol.client.update_goal_state(force_update=True)
+            protocol.client.reset_goal_state()
             extensions_on_hold = protocol.get_goal_state().extensions_goal_state.on_hold
             self.assertFalse(extensions_on_hold, "Extensions should not be on hold when the in-vm artifacts profile response body is None")
 
             mock_response = MockHttpResponse(200, '   '.encode('utf-8'))
-            protocol.client.update_goal_state(force_update=True)
+            protocol.client.reset_goal_state()
             extensions_on_hold = protocol.get_goal_state().extensions_goal_state.on_hold
             self.assertFalse(extensions_on_hold, "Extensions should not be on hold when the in-vm artifacts profile response is an empty string")
 
             mock_response = MockHttpResponse(200, '{ }'.encode('utf-8'))
-            protocol.client.update_goal_state(force_update=True)
+            protocol.client.reset_goal_state()
             extensions_on_hold = protocol.get_goal_state().extensions_goal_state.on_hold
             self.assertFalse(extensions_on_hold, "Extensions should not be on hold when the in-vm artifacts profile response is an empty json object")
 
             with patch("azurelinuxagent.common.protocol.extensions_goal_state_from_extensions_config.add_event") as add_event:
                 mock_response = MockHttpResponse(200, 'invalid json'.encode('utf-8'))
-                protocol.client.update_goal_state(force_update=True)
+                protocol.client.reset_goal_state()
 
                 extensions_on_hold = protocol.get_goal_state().extensions_goal_state.on_hold
                 self.assertFalse(extensions_on_hold, "Extensions should not be on hold when the in-vm artifacts profile response is not valid json")
@@ -780,7 +781,7 @@ class TestWireClient(HttpRequestPredicates, AgentTestCase):
             protocol.set_http_handlers(http_get_handler=http_get_handler)
             HostPluginProtocol.is_default_channel = False
 
-            protocol.client.update_goal_state(force_update=True)
+            protocol.client.reset_goal_state()
 
             urls = protocol.get_tracked_urls()
             self.assertEqual(len(urls), 1, "Unexpected HTTP requests: [{0}]".format(urls))
@@ -799,7 +800,7 @@ class TestWireClient(HttpRequestPredicates, AgentTestCase):
 
             HostPluginProtocol.is_default_channel = False
             try:
-                protocol.client.update_goal_state(force_update=True)
+                protocol.client.reset_goal_state()
 
                 urls = protocol.get_tracked_urls()
                 self.assertEqual(len(urls), 2, "Invalid number of requests: [{0}]".format(urls))
@@ -832,7 +833,7 @@ class TestWireClient(HttpRequestPredicates, AgentTestCase):
 
                 protocol.set_http_handlers(http_get_handler=http_get_handler)
 
-                protocol.client.update_goal_state(force_update=True)
+                protocol.client.reset_goal_state()
 
                 urls = protocol.get_tracked_urls()
                 self.assertEqual(len(urls), 4, "Invalid number of requests: [{0}]".format(urls))
@@ -864,7 +865,7 @@ class TestWireClient(HttpRequestPredicates, AgentTestCase):
 
             protocol.set_http_handlers(http_get_handler=http_get_handler)
 
-            protocol.client.update_goal_state(force_update=True)
+            protocol.client.reset_goal_state()
 
             urls = protocol.get_tracked_urls()
             self.assertEqual(len(urls), 4, "Invalid number of requests: [{0}]".format(urls))
@@ -991,7 +992,7 @@ class TestWireClient(HttpRequestPredicates, AgentTestCase):
 
 class UpdateGoalStateTestCase(HttpRequestPredicates, AgentTestCase):
     """
-    Tests for WireClient.update_goal_state()
+    Tests for WireClient.update_goal_state() and WireClient.reset_goal_state()
     """
 
     def test_it_should_update_the_goal_state_and_the_host_plugin_when_the_incarnation_changes(self):
@@ -1036,7 +1037,7 @@ class UpdateGoalStateTestCase(HttpRequestPredicates, AgentTestCase):
                 '''
 
                 if forced:
-                    protocol.client.update_goal_state(force_update=True)
+                    protocol.client.reset_goal_state()
                 else:
                     protocol.client.update_goal_state()
 
@@ -1090,7 +1091,7 @@ class UpdateGoalStateTestCase(HttpRequestPredicates, AgentTestCase):
             protocol.mock_wire_data.set_role_config_name(new_role_config_name)
             protocol.mock_wire_data.shared_config = new_shared_conf
 
-            protocol.client.update_goal_state(force_update=True)
+            protocol.client.reset_goal_state()
 
             self.assertEqual(protocol.client.get_goal_state().incarnation, incarnation)
             self.assertEqual(protocol.client.get_shared_conf().xml_text, new_shared_conf)
@@ -1098,6 +1099,28 @@ class UpdateGoalStateTestCase(HttpRequestPredicates, AgentTestCase):
             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_reset_should_init_provided_goal_state_properties(self):
+        with mock_wire_protocol(mockwiredata.DATA_FILE) as protocol:
+            protocol.client.reset_goal_state(goal_state_properties=GoalStateProperties.All & ~GoalStateProperties.ExtensionsGoalState)
+
+            with self.assertRaises(ProtocolError) as context:
+                _ = protocol.client.get_certs()
+
+            expected_message = "Certificates is not in goal state properties"
+            self.assertIn(expected_message, str(context.exception))
+
+    def test_reset_should_init_the_goal_state(self):
+        with mock_wire_protocol(mockwiredata.DATA_FILE) as protocol:
+            new_container_id = str(uuid.uuid4())
+            new_role_config_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.client.reset_goal_state()
+
+            self.assertEqual(protocol.client.get_goal_state().container_id, new_container_id)
+            self.assertEqual(protocol.client.get_goal_state().role_config_name, new_role_config_name)
+
 
 class UpdateHostPluginFromGoalStateTestCase(AgentTestCase):
     """