address comments

azurelinuxagent/ga/agent_update_handler.py

@@ -25,6 +25,7 @@
 from azurelinuxagent.common.utils import textutil
 from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
 from azurelinuxagent.common.version import get_daemon_version
+from azurelinuxagent.ga.guestagent import GuestAgentUpdateUtil
 from azurelinuxagent.ga.rsm_version_updater import RSMVersionUpdater
 from azurelinuxagent.ga.self_update_version_updater import SelfUpdateVersionUpdater
 
@@ -41,10 +42,6 @@ def get_agent_update_handler(protocol):
     return AgentUpdateHandler(protocol)
 
 
-RSM_UPDATE_STATE_FILE = "waagent_rsm_update"
-INITIAL_UPDATE_STATE_FILE = "waagent_initial_update"
-
-
 class AgentUpdateHandler(object):
     """
     This class handles two type of agent updates. Handler initializes the updater to SelfUpdateVersionUpdater and switch to appropriate updater based on below conditions:
@@ -84,7 +81,7 @@ def __init__(self, protocol):
         self._last_attempted_update_error_msg = ""
 
         # Restore the state of rsm update. Default to self-update if last update is not with RSM or if agent doing initial update
-        if not self._get_is_last_update_with_rsm() or self._is_initial_update():
+        if not GuestAgentUpdateUtil.is_last_update_with_rsm() or GuestAgentUpdateUtil.is_initial_update():
             self._updater = SelfUpdateVersionUpdater(self._gs_id)
         else:
             self._updater = RSMVersionUpdater(self._gs_id, self._daemon_version)
@@ -98,68 +95,6 @@ def _get_daemon_version_for_update():
         # use the min version as 2.2.53 as we started setting the daemon version starting 2.2.53.
         return FlexibleVersion("2.2.53")
 
-    @staticmethod
-    def _get_initial_update_state_file():
-        """
-        This file keeps if initial update is attempted or not
-        """
-        return os.path.join(conf.get_lib_dir(), INITIAL_UPDATE_STATE_FILE)
-
-    def _save_initial_update_state_file(self):
-        """
-        Save the file if agent attempted initial update
-        """
-        try:
-            with open(self._get_initial_update_state_file(), "w"):
-                pass
-        except Exception as e:
-            msg = "Error creating the initial update state file ({0}): {1}".format(self._get_initial_update_state_file(), ustr(e))
-            logger.warn(msg)
-            add_event(op=WALAEventOperation.AgentUpgrade, message=msg, log_event=False)
-
-    def _is_initial_update(self):
-        """
-        Returns True if state file doesn't exit as presence of file consider as initial update already attempted
-        """
-        return not os.path.exists(self._get_initial_update_state_file())
-
-    @staticmethod
-    def _get_rsm_update_state_file():
-        """
-        This file keeps if last attempted update is rsm or not.
-        """
-        return os.path.join(conf.get_lib_dir(), RSM_UPDATE_STATE_FILE)
-
-    def _save_rsm_update_state_file(self):
-        """
-        Save the rsm state empty file when we switch to RSM
-        """
-        try:
-            with open(self._get_rsm_update_state_file(), "w"):
-                pass
-        except Exception as e:
-            msg = "Error creating the RSM state file ({0}): {1}".format(self._get_rsm_update_state_file(), ustr(e))
-            logger.warn(msg)
-            add_event(op=WALAEventOperation.AgentUpgrade, message=msg, log_event=False)
-
-    def _remove_rsm_update_state_file(self):
-        """
-        Remove the rsm state file when we switch to self-update
-        """
-        try:
-            if os.path.exists(self._get_rsm_update_state_file()):
-                os.remove(self._get_rsm_update_state_file())
-        except Exception as e:
-            msg = "Error removing the RSM state file ({0}): {1}".format(self._get_rsm_update_state_file(), ustr(e))
-            logger.warn(msg)
-            add_event(op=WALAEventOperation.AgentUpgrade, message=msg, log_event=False)
-
-    def _get_is_last_update_with_rsm(self):
-        """
-        Returns True if state file exists as this consider as last update with RSM is true
-        """
-        return os.path.exists(self._get_rsm_update_state_file())
-
     def _get_agent_family_manifest(self, goal_state):
         """
         Get the agent_family from last GS for the given family
@@ -214,9 +149,7 @@ def run(self, goal_state, ext_gs_updated):
 
             # Always agent uses self-update for initial update regardless vm enrolled into RSM or not
             # So ignoring the check for updater switch for the initial goal state/update
-            initial_update = self._is_initial_update()
-            if not initial_update:
-
+            if not GuestAgentUpdateUtil.is_initial_update():
                 # Updater will return True or False if we need to switch the updater
                 # If self-updater receives RSM update enabled, it will switch to RSM updater
                 # If RSM updater receives RSM update disabled, it will switch to self-update
@@ -228,14 +161,14 @@ def run(self, goal_state, ext_gs_updated):
                     logger.info(msg)
                     add_event(op=WALAEventOperation.AgentUpgrade, message=msg, log_event=False)
                     self._updater = SelfUpdateVersionUpdater(self._gs_id)
-                    self._remove_rsm_update_state_file()
+                    GuestAgentUpdateUtil.remove_rsm_update_state_file()
 
                 if is_rsm_update_enabled and isinstance(self._updater, SelfUpdateVersionUpdater):
                     msg = "VM enabled for RSM updates, switching to RSM update mode"
                     logger.info(msg)
                     add_event(op=WALAEventOperation.AgentUpgrade, message=msg, log_event=False)
                     self._updater = RSMVersionUpdater(self._gs_id, self._daemon_version)
-                    self._save_rsm_update_state_file()
+                    GuestAgentUpdateUtil.save_rsm_update_state_file()
 
             # If updater is changed in previous step, we allow update as it consider as first attempt. If not, it checks below condition
             # RSM checks new goal state; self-update checks manifest download interval
@@ -244,7 +177,7 @@ def run(self, goal_state, ext_gs_updated):
 
             self._updater.retrieve_agent_version(agent_family, goal_state)
 
-            if not self._updater.is_retrieved_version_allowed_to_update(agent_family, initial_update):
+            if not self._updater.is_retrieved_version_allowed_to_update(agent_family):
                 return
             self._updater.log_new_agent_update_message()
             agent = self._updater.download_and_get_new_agent(self._protocol, agent_family, goal_state)
@@ -285,8 +218,8 @@ def run(self, goal_state, ext_gs_updated):
 
         # save initial update state when agent is doing first update
         finally:
-            if self._is_initial_update():
-                self._save_initial_update_state_file()
+            if GuestAgentUpdateUtil.is_initial_update():
+                GuestAgentUpdateUtil.save_initial_update_state_file()
 
     def get_vmagent_update_status(self):
         """
@@ -308,4 +241,4 @@ def get_vmagent_update_status(self):
             msg = "Unable to report agent update status: {0}".format(textutil.format_exception(err))
             logger.warn(msg)
             add_event(op=WALAEventOperation.AgentUpgrade, is_success=False, message=msg, log_event=True)
-        return None
+        return None

azurelinuxagent/ga/ga_version_updater.py

@@ -63,11 +63,10 @@ def retrieve_agent_version(self, agent_family, goal_state):
         """
         raise NotImplementedError
 
-    def is_retrieved_version_allowed_to_update(self, agent_family, initial_update):
+    def is_retrieved_version_allowed_to_update(self, agent_family):
         """
         Checks all base condition if new version allow to update.
         @param agent_family: agent family
-        @param initial_update: True if it's initial update else False
         @return: True if allowed to update else False
         """
         raise NotImplementedError

azurelinuxagent/ga/guestagent.py

@@ -18,6 +18,9 @@
 MAX_FAILURE = 3  # Max failure allowed for agent before declare bad agent
 AGENT_UPDATE_COUNT_FILE = "update_attempt.json"  # File for tracking agent update attempt count
 
+RSM_UPDATE_STATE_FILE = "waagent_rsm_update"
+INITIAL_UPDATE_STATE_FILE = "waagent_initial_update"
+
 
 class GuestAgent(object):
     def __init__(self, path, pkg):
@@ -329,3 +332,71 @@ def to_json(self):
         return data
 
 
+class GuestAgentUpdateUtil(object):
+
+    @staticmethod
+    def get_initial_update_state_file():
+        """
+        This file tracks whether the initial update attempt has been made or not
+        """
+        return os.path.join(conf.get_lib_dir(), INITIAL_UPDATE_STATE_FILE)
+
+    @staticmethod
+    def save_initial_update_state_file():
+        """
+        Save the file if agent attempted initial update
+        """
+        try:
+            with open(GuestAgentUpdateUtil.get_initial_update_state_file(), "w"):
+                pass
+        except Exception as e:
+            msg = "Error creating the initial update state file ({0}): {1}".format(GuestAgentUpdateUtil.get_initial_update_state_file(), ustr(e))
+            logger.warn(msg)
+            add_event(op=WALAEventOperation.AgentUpgrade, message=msg, log_event=False)
+
+    @staticmethod
+    def is_initial_update():
+        """
+        Returns True if the state file doesn't exist, as the presence of the file indicates that the initial update has already been attempted
+        """
+        return not os.path.exists(GuestAgentUpdateUtil.get_initial_update_state_file())
+
+    @staticmethod
+    def get_rsm_update_state_file():
+        """
+        This file tracks whether the last attempted update was an RSM update or not
+        """
+        return os.path.join(conf.get_lib_dir(), RSM_UPDATE_STATE_FILE)
+
+    @staticmethod
+    def save_rsm_update_state_file():
+        """
+        Save the rsm state empty file when we switch to RSM
+        """
+        try:
+            with open(GuestAgentUpdateUtil.get_rsm_update_state_file(), "w"):
+                pass
+        except Exception as e:
+            msg = "Error creating the RSM state file ({0}): {1}".format(GuestAgentUpdateUtil.get_rsm_update_state_file(), ustr(e))
+            logger.warn(msg)
+            add_event(op=WALAEventOperation.AgentUpgrade, message=msg, log_event=False)
+
+    @staticmethod
+    def remove_rsm_update_state_file():
+        """
+        Remove the rsm state file when we switch to self-update
+        """
+        try:
+            if os.path.exists(GuestAgentUpdateUtil.get_rsm_update_state_file()):
+                os.remove(GuestAgentUpdateUtil.get_rsm_update_state_file())
+        except Exception as e:
+            msg = "Error removing the RSM state file ({0}): {1}".format(GuestAgentUpdateUtil.get_rsm_update_state_file(), ustr(e))
+            logger.warn(msg)
+            add_event(op=WALAEventOperation.AgentUpgrade, message=msg, log_event=False)
+
+    @staticmethod
+    def is_last_update_with_rsm():
+        """
+        Returns True if the state file exists, as this indicates that the last update was with RSM
+        """
+        return os.path.exists(GuestAgentUpdateUtil.get_rsm_update_state_file())

azurelinuxagent/ga/rsm_version_updater.py

@@ -84,7 +84,7 @@ def retrieve_agent_version(self, agent_family, goal_state):
         """
         self._version = FlexibleVersion(agent_family.version)
 
-    def is_retrieved_version_allowed_to_update(self, agent_family, initial_update):
+    def is_retrieved_version_allowed_to_update(self, agent_family):
         """
         Once version retrieved from goal state, we check if we allowed to update for that version
         allow update If new version not same as current version, not below than daemon version and if version is from rsm request

azurelinuxagent/ga/self_update_version_updater.py

@@ -25,6 +25,7 @@
 from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
 from azurelinuxagent.common.version import CURRENT_VERSION
 from azurelinuxagent.ga.ga_version_updater import GAVersionUpdater
+from azurelinuxagent.ga.guestagent import GuestAgentUpdateUtil
 
 
 class SelfUpdateType(object):
@@ -150,7 +151,7 @@ def retrieve_agent_version(self, agent_family, goal_state):
         largest_version = self._get_largest_version(self._agent_manifest)
         self._version = largest_version
 
-    def is_retrieved_version_allowed_to_update(self, agent_family, initial_update):
+    def is_retrieved_version_allowed_to_update(self, agent_family):
         """
         we don't allow new version update, if
             1) The version is not greater than current version
@@ -164,7 +165,7 @@ def is_retrieved_version_allowed_to_update(self, agent_family, initial_update):
             return False
 
         # very first update need to proceed without any delay
-        if initial_update:
+        if GuestAgentUpdateUtil.is_initial_update():
             return True
 
         if not self._new_agent_allowed_now_to_update():

tests/ga/test_agent_update_handler.py

@@ -12,9 +12,8 @@
 
 from azurelinuxagent.common.protocol.util import ProtocolUtil
 from azurelinuxagent.common.version import CURRENT_VERSION, AGENT_NAME
-from azurelinuxagent.ga.agent_update_handler import get_agent_update_handler, INITIAL_UPDATE_STATE_FILE, \
-    RSM_UPDATE_STATE_FILE
-from azurelinuxagent.ga.guestagent import GuestAgent
+from azurelinuxagent.ga.agent_update_handler import get_agent_update_handler
+from azurelinuxagent.ga.guestagent import GuestAgent, INITIAL_UPDATE_STATE_FILE, RSM_UPDATE_STATE_FILE
 from tests.ga.test_update import UpdateTestCase
 from tests.lib.http_request_predicates import HttpRequestPredicates
 from tests.lib.mock_wire_protocol import mock_wire_protocol, MockHttpResponse

tests/ga/test_update.py

@@ -21,8 +21,7 @@
 from datetime import datetime, timedelta
 from threading import current_thread
 
-from azurelinuxagent.ga.agent_update_handler import INITIAL_UPDATE_STATE_FILE
-from azurelinuxagent.ga.guestagent import GuestAgent, GuestAgentError, AGENT_ERROR_FILE
+from azurelinuxagent.ga.guestagent import GuestAgent, GuestAgentError, AGENT_ERROR_FILE, INITIAL_UPDATE_STATE_FILE
 from azurelinuxagent.common import conf
 from azurelinuxagent.common.event import EVENTS_DIRECTORY, WALAEventOperation
 from azurelinuxagent.common.exception import HttpError, \