Repository: Azure/WALinuxAgent
Pull request #3275 (2024-12-10 01:29:41)

picks the random time for attempting new update 
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description

Today we count self-update time 6/24 hours period as a sliding window from service start. This strategy may not be good for vmss, for example. All the vms in vmss may get update around same time

New strategy picks the random time[0..6/24] when we detect the new update

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
- [ ] Ensure development PR is based on the `develop` branch.
- [ ] 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).


Patch:
diff --git a/azurelinuxagent/ga/agent_update_handler.py b/azurelinuxagent/ga/agent_update_handler.py
index ee6a44f9..30940f30 100644
--- a/azurelinuxagent/ga/agent_update_handler.py
+++ b/azurelinuxagent/ga/agent_update_handler.py
@@ -15,7 +15,6 @@
 # limitations under the License.
 #
 # Requires Python 2.6+ and Openssl 1.0+
-import os
 
 from azurelinuxagent.common import conf, logger
 from azurelinuxagent.common.event import add_event, WALAEventOperation
@@ -25,6 +24,7 @@ from azurelinuxagent.common.protocol.restapi import VMAgentUpdateStatuses, VMAge
 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 +41,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 +80,7 @@ class AgentUpdateHandler(object):
         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 +94,6 @@ class AgentUpdateHandler(object):
         # 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,8 +148,7 @@ class AgentUpdateHandler(object):
 
             # 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
-            if not self._is_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
@@ -227,14 +160,14 @@ class AgentUpdateHandler(object):
                     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
@@ -284,8 +217,8 @@ class AgentUpdateHandler(object):
 
         # 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):
         """
@@ -307,4 +240,4 @@ class AgentUpdateHandler(object):
             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
\ No newline at end of file
diff --git a/azurelinuxagent/ga/guestagent.py b/azurelinuxagent/ga/guestagent.py
index b4b2d05b..3166d8ea 100644
--- a/azurelinuxagent/ga/guestagent.py
+++ b/azurelinuxagent/ga/guestagent.py
@@ -18,6 +18,9 @@ AGENT_MANIFEST_FILE = "HandlerManifest.json"
 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 @@ class GuestAgentUpdateAttempt(object):
         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())
diff --git a/azurelinuxagent/ga/self_update_version_updater.py b/azurelinuxagent/ga/self_update_version_updater.py
index 5a839851..9112453f 100644
--- a/azurelinuxagent/ga/self_update_version_updater.py
+++ b/azurelinuxagent/ga/self_update_version_updater.py
@@ -17,6 +17,7 @@
 # Requires Python 2.6+ and Openssl 1.0+
 
 import datetime
+import random
 
 from azurelinuxagent.common import conf, logger
 from azurelinuxagent.common.event import add_event, WALAEventOperation
@@ -24,6 +25,7 @@ from azurelinuxagent.common.exception import AgentUpgradeExitException, AgentUpd
 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):
@@ -38,7 +40,7 @@ class SelfUpdateVersionUpdater(GAVersionUpdater):
     def __init__(self, gs_id):
         super(SelfUpdateVersionUpdater, self).__init__(gs_id)
         self._last_attempted_manifest_download_time = datetime.datetime.min
-        self._last_attempted_self_update_time = datetime.datetime.min
+        self._next_update_time = datetime.datetime.min
 
     @staticmethod
     def _get_largest_version(agent_manifest):
@@ -61,34 +63,35 @@ class SelfUpdateVersionUpdater(GAVersionUpdater):
         return SelfUpdateType.Regular
 
     @staticmethod
-    def _get_next_process_time(last_val, frequency, now):
+    def _get_next_process_time(upgrade_type, now):
         """
-        Get the next upgrade time
+        Returns random time in between 0 to 24hrs(regular) or 4hrs(hotfix) from now
         """
-        return now if last_val == datetime.datetime.min else last_val + datetime.timedelta(seconds=frequency)
+        if upgrade_type == SelfUpdateType.Hotfix:
+            frequency = conf.get_self_update_hotfix_frequency()
+        else:
+            frequency = conf.get_self_update_regular_frequency()
+        return now + datetime.timedelta(seconds=random.randint(0, frequency))
 
-    def _is_new_agent_allowed_update(self):
+    def _new_agent_allowed_now_to_update(self):
         """
-        This method ensure that update is allowed only once per (hotfix/Regular) upgrade frequency
+        This method is called when a new update is detected and computes a random time for the next update on the first call.
+        Since the method is called periodically until we reach the next update time, we shouldn't refresh or recompute the next update time on every call.
+        We use default value(datetime.datetime.min) to ensure the computation happens only once. This next_update_time will reset to default value(datetime.min) when agent allowed to update.
+        So that, in case the update fails due to an issue, such as a package download error, the same default value used to recompute the next update time.
         """
         now = datetime.datetime.utcnow()
         upgrade_type = self._get_agent_upgrade_type(self._version)
-        if upgrade_type == SelfUpdateType.Hotfix:
-            next_update_time = self._get_next_process_time(self._last_attempted_self_update_time,
-                                                           conf.get_self_update_hotfix_frequency(), now)
-        else:
-            next_update_time = self._get_next_process_time(self._last_attempted_self_update_time,
-                                                           conf.get_self_update_regular_frequency(), now)
 
-        if self._version > CURRENT_VERSION:
-            message = "Self-update discovered new {0} upgrade WALinuxAgent-{1}; Will upgrade on or after {2}".format(
-                upgrade_type, str(self._version), next_update_time.strftime(logger.Logger.LogTimeFormatInUTC))
-            logger.info(message)
-            add_event(op=WALAEventOperation.AgentUpgrade, message=message, log_event=False)
-
-        if next_update_time <= now:
-            # Update the last upgrade check time even if no new agent is available for upgrade
-            self._last_attempted_self_update_time = now
+        if self._next_update_time == datetime.datetime.min:
+            self._next_update_time = self._get_next_process_time(upgrade_type, now)
+        message = "Self-update discovered new {0} upgrade WALinuxAgent-{1}; Will upgrade on or after {2}".format(
+            upgrade_type, str(self._version), self._next_update_time.strftime(logger.Logger.LogTimeFormatInUTC))
+        logger.info(message)
+        add_event(op=WALAEventOperation.AgentUpgrade, message=message, log_event=False)
+
+        if self._next_update_time <= now:
+            self._next_update_time = datetime.datetime.min
             return True
         return False
 
@@ -150,14 +153,22 @@ class SelfUpdateVersionUpdater(GAVersionUpdater):
 
     def is_retrieved_version_allowed_to_update(self, agent_family):
         """
-        checks update is spread per (as specified in the conf.get_self_update_hotfix_frequency() or conf.get_self_update_regular_frequency())
-        or if version below than current version
-        return false when we don't allow updates.
+        we don't allow new version update, if
+            1) The version is not greater than current version
+            2) if current time is before next update time
+
+        Allow the update, if
+            1) Initial update
+            2) If current time is on or after next update time
         """
-        if not self._is_new_agent_allowed_update():
+        if self._version <= CURRENT_VERSION:
             return False
 
-        if self._version <= CURRENT_VERSION:
+        # very first update need to proceed without any delay
+        if GuestAgentUpdateUtil.is_initial_update():
+            return True
+
+        if not self._new_agent_allowed_now_to_update():
             return False
 
         return True


Test patch:
diff --git a/tests/ga/test_agent_update_handler.py b/tests/ga/test_agent_update_handler.py
index c2d01a42..02fadb46 100644
--- a/tests/ga/test_agent_update_handler.py
+++ b/tests/ga/test_agent_update_handler.py
@@ -1,6 +1,8 @@
 import contextlib
 import json
 import os
+import random
+import time
 
 from azurelinuxagent.common import conf
 from azurelinuxagent.common.event import WALAEventOperation
@@ -10,9 +12,8 @@ from azurelinuxagent.common.protocol.restapi import VMAgentUpdateStatuses
 
 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
@@ -29,7 +30,7 @@ class TestAgentUpdate(UpdateTestCase):
         clear_singleton_instances(ProtocolUtil)
 
     @contextlib.contextmanager
-    def _get_agent_update_handler(self, test_data=None, autoupdate_frequency=0.001, autoupdate_enabled=True, initial_update_attempted=True, protocol_get_error=False, mock_get_header=None, mock_put_header=None):
+    def _get_agent_update_handler(self, test_data=None, autoupdate_frequency=0.001, autoupdate_enabled=True, initial_update_attempted=True, protocol_get_error=False, mock_get_header=None, mock_put_header=None, mock_random_update_time=True):
         # Default to DATA_FILE of test_data parameter raises the pylint warning
         # W0102: Dangerous default value DATA_FILE (builtins.dict) as argument (dangerous-default-value)
         test_data = DATA_FILE if test_data is None else test_data
@@ -61,14 +62,25 @@ class TestAgentUpdate(UpdateTestCase):
             if initial_update_attempted:
                 open(os.path.join(conf.get_lib_dir(), INITIAL_UPDATE_STATE_FILE), "a").close()
 
+            original_randint = random.randint
+
+            def _mock_random_update_time(a, b):
+                if mock_random_update_time:
+                    return 0
+                # some tests mock normal/hotfix frequency to 1 second, return 0.001 to avoid long delay and still test the logic
+                if b == 1:
+                    return 0.001
+                return original_randint(a, b)
+
             with patch("azurelinuxagent.common.conf.get_autoupdate_enabled", return_value=autoupdate_enabled):
                 with patch("azurelinuxagent.common.conf.get_autoupdate_frequency", return_value=autoupdate_frequency):
-                    with patch("azurelinuxagent.common.conf.get_autoupdate_gafamily", return_value="Prod"):
-                        with patch("azurelinuxagent.common.conf.get_enable_ga_versioning", return_value=True):
-                            with patch("azurelinuxagent.common.event.EventLogger.add_event") as mock_telemetry:
-                                agent_update_handler = get_agent_update_handler(protocol)
-                                agent_update_handler._protocol = protocol
-                                yield agent_update_handler, mock_telemetry
+                    with patch("azurelinuxagent.ga.self_update_version_updater.random.randint", side_effect=_mock_random_update_time):
+                        with patch("azurelinuxagent.common.conf.get_autoupdate_gafamily", return_value="Prod"):
+                            with patch("azurelinuxagent.common.conf.get_enable_ga_versioning", return_value=True):
+                                with patch("azurelinuxagent.common.event.EventLogger.add_event") as mock_telemetry:
+                                    agent_update_handler = get_agent_update_handler(protocol)
+                                    agent_update_handler._protocol = protocol
+                                    yield agent_update_handler, mock_telemetry
 
     def _assert_agent_directories_available(self, versions):
         for version in versions:
@@ -130,7 +142,7 @@ class TestAgentUpdate(UpdateTestCase):
             self._assert_agent_directories_exist_and_others_dont_exist(versions=[str(CURRENT_VERSION), "99999.0.0.0"])
             self._assert_agent_exit_process_telemetry_emitted(ustr(context.exception.reason))
 
-    def test_it_should_not_update_to_largest_version_if_time_window_not_elapsed(self):
+    def test_it_should_not_update_to_largest_version_if_manifest_download_time_not_elapsed(self):
         self.prepare_agents(count=1)
 
         data_file = DATA_FILE.copy()
@@ -146,14 +158,30 @@ class TestAgentUpdate(UpdateTestCase):
             self.assertFalse(os.path.exists(self.agent_dir("99999.0.0.0")),
                              "New agent directory should not be found")
 
-    def test_it_should_update_to_largest_version_if_time_window_elapsed(self):
+    def test_it_should_not_do_self_update_if_update_time_is_not_elapsed(self):
+        self.prepare_agents(count=1)
+
+        data_file = DATA_FILE.copy()
+        data_file["ga_manifest"] = "wire/ga_manifest_no_uris.xml"
+        with self._get_agent_update_handler(test_data=data_file, mock_random_update_time=False) as (agent_update_handler, _):
+            agent_update_handler.run(agent_update_handler._protocol.get_goal_state(), True)
+            self.assertFalse(os.path.exists(self.agent_dir("99999.0.0.0")),
+                             "New agent directory should not be found")
+            agent_update_handler._protocol.mock_wire_data.set_ga_manifest("wire/ga_manifest.xml")
+            agent_update_handler._protocol.mock_wire_data.set_incarnation(2)
+            agent_update_handler._protocol.client.update_goal_state()
+            agent_update_handler.run(agent_update_handler._protocol.get_goal_state(), True)
+            self.assertFalse(os.path.exists(self.agent_dir("99999.0.0.0")),
+                             "New agent directory should not be found")
+
+    def test_it_should_update_to_largest_version_after_time_window_elapsed(self):
         self.prepare_agents(count=1)
 
         data_file = DATA_FILE.copy()
         data_file["ga_manifest"] = "wire/ga_manifest_no_uris.xml"
-        with patch("azurelinuxagent.common.conf.get_self_update_hotfix_frequency", return_value=0.001):
-            with patch("azurelinuxagent.common.conf.get_self_update_regular_frequency", return_value=0.001):
-                with self._get_agent_update_handler(test_data=data_file) as (agent_update_handler, mock_telemetry):
+        with patch("azurelinuxagent.common.conf.get_self_update_hotfix_frequency", return_value=1):
+            with patch("azurelinuxagent.common.conf.get_self_update_regular_frequency", return_value=1):
+                with self._get_agent_update_handler(test_data=data_file, mock_random_update_time=False) as (agent_update_handler, mock_telemetry):
                     with self.assertRaises(AgentUpgradeExitException) as context:
                         agent_update_handler.run(agent_update_handler._protocol.get_goal_state(), True)
                         self.assertFalse(os.path.exists(self.agent_dir("99999.0.0.0")),
@@ -161,6 +189,8 @@ class TestAgentUpdate(UpdateTestCase):
                         agent_update_handler._protocol.mock_wire_data.set_ga_manifest("wire/ga_manifest.xml")
                         agent_update_handler._protocol.mock_wire_data.set_incarnation(2)
                         agent_update_handler._protocol.client.update_goal_state()
+                        # sleeping for update window to elapse
+                        time.sleep(0.1)
                         agent_update_handler.run(agent_update_handler._protocol.get_goal_state(), True)
                     self._assert_update_discovered_from_agent_manifest(mock_telemetry, inc=2, version="99999.0.0.0")
                     self._assert_agent_directories_exist_and_others_dont_exist(versions=[str(CURRENT_VERSION), "99999.0.0.0"])
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index 89e38e6d..e97ca59a 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -7,6 +7,7 @@ import contextlib
 import glob
 import json
 import os
+import random
 import re
 import shutil
 import stat
@@ -20,8 +21,7 @@ import zipfile
 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.logger import LogLevel
 from azurelinuxagent.common.event import EVENTS_DIRECTORY, WALAEventOperation
@@ -1327,7 +1327,7 @@ class TestAgentUpgrade(UpdateTestCase):
 
     @contextlib.contextmanager
     def __get_update_handler(self, iterations=1, test_data=None,
-                             reload_conf=None, autoupdate_frequency=0.001, hotfix_frequency=1.0, normal_frequency=2.0, initial_update_attempted=True):
+                             reload_conf=None, autoupdate_frequency=0.001, hotfix_frequency=10, normal_frequency=10, initial_update_attempted=True, mock_random_update_time=True):
 
         if initial_update_attempted:
             open(os.path.join(conf.get_lib_dir(), INITIAL_UPDATE_STATE_FILE), "a").close()
@@ -1355,11 +1355,22 @@ class TestAgentUpgrade(UpdateTestCase):
                 protocol.aggregate_status = json.loads(args[0])
                 return MockHttpResponse(status=201)
 
+            original_randint = random.randint
+
+            def _mock_random_update_time(a, b):
+                if mock_random_update_time:
+                    return 0
+                if b == 1:  # some tests mock normal/hotfix frequency to 1 second, return 0.001 to avoid long delay and still test the logic
+                    return 0.001
+                return original_randint(a, b)
+
             protocol.set_http_handlers(http_get_handler=get_handler, http_put_handler=put_handler)
             with self.create_conf_mocks(autoupdate_frequency, hotfix_frequency, normal_frequency):
-                with patch("azurelinuxagent.common.event.EventLogger.add_event") as mock_telemetry:
-                    update_handler._protocol = protocol
-                    yield update_handler, mock_telemetry
+                with patch("azurelinuxagent.ga.self_update_version_updater.random.randint",
+                           side_effect=_mock_random_update_time):
+                    with patch("azurelinuxagent.common.event.EventLogger.add_event") as mock_telemetry:
+                        update_handler._protocol = protocol
+                        yield update_handler, mock_telemetry
 
     def __assert_exit_code_successful(self, update_handler):
         self.assertEqual(0, update_handler.get_exit_code(), "Exit code should be 0")
@@ -1606,8 +1617,7 @@ class TestAgentUpgrade(UpdateTestCase):
         data_file = wire_protocol_data.DATA_FILE.copy()
         # This is to fail the agent update at first attempt so that agent doesn't go through update
         data_file["ga_manifest"] = "wire/ga_manifest_no_uris.xml"
-        with self.__get_update_handler(iterations=no_of_iterations, test_data=data_file, reload_conf=reload_conf,
-                                       hotfix_frequency=10, normal_frequency=10) as (update_handler, _):
+        with self.__get_update_handler(iterations=no_of_iterations, test_data=data_file, reload_conf=reload_conf, mock_random_update_time=False) as (update_handler, _):
             update_handler._protocol.mock_wire_data.set_incarnation(2)
             update_handler.run(debug=True)
 
@@ -1646,7 +1656,7 @@ class TestAgentUpgrade(UpdateTestCase):
         data_file = wire_protocol_data.DATA_FILE.copy()
         data_file["ga_manifest"] = "wire/ga_manifest_no_uris.xml"
         with self.__get_update_handler(iterations=no_of_iterations, test_data=data_file, reload_conf=reload_conf,
-                                       hotfix_frequency=0.001, normal_frequency=0.001) as (update_handler, mock_telemetry):
+                                       hotfix_frequency=1, normal_frequency=1, mock_random_update_time=False) as (update_handler, mock_telemetry):
             update_handler._protocol.mock_wire_data.set_incarnation(2)
             update_handler.run(debug=True)