fix agent update UT (Azure#3051)

(cherry picked from commit d9f7ed4)
nagworld9 · Feb 14, 2024 · 75843d2 · 75843d2
1 parent b498fd5
commit 75843d2
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/tests/ga/test_agent_update_handler.py b/tests/ga/test_agent_update_handler.py
@@ -494,6 +494,7 @@ def test_it_should_continue_with_update_if_number_of_update_attempts_less_than_3
                 with self.assertRaises(AgentUpgradeExitException):
                     agent_update_handler._protocol.mock_wire_data.set_version_in_agent_family(
                         str(latest_version))
+                    agent_update_handler._protocol.mock_wire_data.set_version_in_ga_manifest(str(latest_version))
                     agent_update_handler._protocol.mock_wire_data.set_incarnation(i+2)
                     agent_update_handler._protocol.client.update_goal_state()
                     agent_update_handler.run(agent_update_handler._protocol.get_goal_state(), True)

diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
@@ -1849,7 +1849,7 @@ def reload_conf(url, protocol):
         data_file['ga_manifest'] = "wire/ga_manifest_no_upgrade.xml"
         # Setting the prod frequency to mimic a real scenario
         with self.__get_update_handler(iterations=no_of_iterations, test_data=data_file, reload_conf=reload_conf, autoupdate_frequency=6000) as (update_handler, mock_telemetry):
-            update_handler._protocol.mock_wire_data.set_ga_manifest_version_version(str(CURRENT_VERSION))
+            update_handler._protocol.mock_wire_data.set_version_in_ga_manifest(str(CURRENT_VERSION))
             update_handler._protocol.mock_wire_data.set_incarnation(20)
             update_handler.run(debug=True)
 

diff --git a/tests/lib/wire_protocol_data.py b/tests/lib/wire_protocol_data.py
@@ -469,5 +469,5 @@ def set_version_in_agent_family(self, version):
     def set_extension_config_is_vm_enabled_for_rsm_upgrades(self, is_vm_enabled_for_rsm_upgrades):
         self.ext_conf = WireProtocolData.replace_xml_element_value(self.ext_conf, "IsVMEnabledForRSMUpgrades", is_vm_enabled_for_rsm_upgrades)
 
-    def set_ga_manifest_version_version(self, version):
+    def set_version_in_ga_manifest(self, version):
         self.ga_manifest = WireProtocolData.replace_xml_element_value(self.ga_manifest, "Version", version)