[AKS] az aks update: Make specified version to match current version when turning off autoupgrade

src/azure-cli/azure/cli/command_modules/acs/custom.py

@@ -843,6 +843,8 @@ def aks_update(
     try:
         # update mc profile
         mc = aks_update_decorator.update_mc_profile_default()
+        # Update kubernetes_version and orchestrator_version based on auto_upgrade_channel
+        mc = aks_update_decorator.update_kubernetes_version_and_orchestrator_version(mc)
     except DecoratorEarlyExitException:
         # exit gracefully
         return None

src/azure-cli/azure/cli/command_modules/acs/managed_cluster_decorator.py

@@ -8484,6 +8484,36 @@ def update_mc_profile_default(self) -> ManagedCluster:
         mc = self.update_bootstrap_profile(mc)
         return mc
 
+    def update_kubernetes_version_and_orchestrator_version(self, mc: ManagedCluster) -> ManagedCluster:
+        """Update kubernetes version and orchestrator version for the ManagedCluster object.
+
+        :param mc: The ManagedCluster object to be updated.
+        :return: The updated ManagedCluster object.
+        """
+        self._ensure_mc(mc)
+
+        # Check if auto_upgrade_channel is set to "none"
+        auto_upgrade_channel = self.context.get_auto_upgrade_channel()
+        if auto_upgrade_channel == "none":
+            warning_message = (
+                "Since auto-upgrade-channel is set to none, cluster kubernetesVersion will be set to the value of "
+                "currentKubernetesVersion, all agent pools orchestratorVersion will be set to the value of "
+                "currentOrchestratorVersion respectively. Continue?"
+            )
+            if not prompt_y_n(warning_message, default="n"):
+                raise CLIError("Operation cancelled by user.")
+
+            # Set kubernetes version to match the current kubernetes version if it has a value
+            if mc.current_kubernetes_version:
+                mc.kubernetes_version = mc.current_kubernetes_version
+
+            # Set orchestrator version for each agent pool to match the current orchestrator version if it has a value
+            for agent_pool in mc.agent_pool_profiles:
+                if agent_pool.current_orchestrator_version:
+                    agent_pool.orchestrator_version = agent_pool.current_orchestrator_version
+
+        return mc
+
     def check_is_postprocessing_required(self, mc: ManagedCluster) -> bool:
         """Helper function to check if postprocessing is required after sending a PUT request to create the cluster.
 

src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_custom.py

@@ -22,6 +22,7 @@
     _get_command_context,
     _update_addons,
     aks_stop,
+    aks_update,
     k8s_install_kubectl,
     k8s_install_kubelogin,
     merge_kubernetes_configurations,
@@ -820,6 +821,43 @@ def test_aks_stop(self):
         )
         self.assertEqual(aks_stop(self.cmd, self.client, "rg", "name", False), None)
 
+    @mock.patch('azure.cli.command_modules.acs.custom.prompt_y_n', return_value=True)
+    @mock.patch('azure.cli.command_modules.acs.custom.get_subscription_id', return_value='00000000-0000-0000-0000-000000000000')
+    @mock.patch('azure.cli.command_modules.acs.custom.cf_agent_pools')
+    @mock.patch('azure.cli.command_modules.acs.managed_cluster_decorator.AKSManagedClusterUpdateDecorator.update_mc_profile_default')
+    @mock.patch('azure.cli.command_modules.acs.managed_cluster_decorator.AKSManagedClusterUpdateDecorator.update_kubernetes_version_and_orchestrator_version')
+    @mock.patch('azure.cli.command_modules.acs.managed_cluster_decorator.AKSManagedClusterUpdateDecorator.update_mc')
+    def test_aks_update_auto_upgrade_channel_none(self, mock_update_mc, mock_update_kubernetes_version_and_orchestrator_version, mock_update_mc_profile_default, mock_cf_agent_pools, mock_get_subscription_id, mock_prompt_y_n):
+        resource_group_name = 'test_rg'
+        name = 'test_cluster'
+        auto_upgrade_channel = 'none'
+
+        # Create a ManagedCluster object
+        mc = mock.Mock()
+        mc.kubernetes_version = '1.18.14'
+        mc.agent_pool_profiles = [
+            mock.Mock(orchestrator_version='1.18.14'),
+            mock.Mock(orchestrator_version='1.19.7')
+        ]
+        mc.current_kubernetes_version = '1.21.2'
+        mc.agent_pool_profiles[0].current_orchestrator_version = '1.21.2'
+        mc.agent_pool_profiles[1].current_orchestrator_version = '1.21.2'
+
+        mock_update_mc_profile_default.return_value = mc
+        mock_update_kubernetes_version_and_orchestrator_version.return_value = mc
+
+        aks_update(
+            cmd=self.cmd,
+            client=self.client,
+            resource_group_name=resource_group_name,
+            name=name,
+            auto_upgrade_channel=auto_upgrade_channel
+        )
+
+        # Check if update_kubernetes_version_and_orchestrator_version was called
+        mock_update_kubernetes_version_and_orchestrator_version.assert_called_once_with(mc)
+        # Check if update_mc was called
+        mock_update_mc.assert_called_once_with(mc)
 
 class TestRunCommand(unittest.TestCase):
     def test_get_command_context_invalid_file(self):

src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_managed_cluster_decorator.py

@@ -11209,6 +11209,123 @@ def test_update_mc_profile_default(self):
 
         dec_1.context.raw_param.print_usage_statistics()
 
+    def test_update_kubernetes_version_and_orchestrator_version(self):
+        # First test case
+        raw_param_dict = {
+            "auto_upgrade_channel": "none",
+        }
+        dec_1 = AKSManagedClusterUpdateDecorator(
+            self.cmd,
+            self.client,
+            raw_param_dict,
+            ResourceType.MGMT_CONTAINERSERVICE,
+        )
+        mc_1 = self.models.ManagedCluster(
+            location="test_location",
+            kubernetes_version="1.18.14",
+            agent_pool_profiles=[
+                self.models.ManagedClusterAgentPoolProfile(
+                    name="agentpool1",
+                    orchestrator_version="1.18.14"
+                ),
+                self.models.ManagedClusterAgentPoolProfile(
+                    name="agentpool2",
+                    orchestrator_version="1.19.7"
+                )
+            ]
+        )
+        # Set current kubernetes version in the managed cluster
+        mc_1.current_kubernetes_version = "1.21.2"
+
+        # Set current orchestrator version in agent pool profiles
+        mc_1.agent_pool_profiles[0].current_orchestrator_version = "1.21.2"
+        mc_1.agent_pool_profiles[1].current_orchestrator_version = "1.21.2"
+
+        dec_1.context.attach_mc(mc_1)
+        with patch(
+            "azure.cli.command_modules.acs.managed_cluster_decorator.prompt_y_n",
+            return_value=True,
+        ):
+            updated_mc_1 = dec_1.update_kubernetes_version_and_orchestrator_version(mc_1)
+
+        # Check if kubernetes_version is updated
+        self.assertEqual(updated_mc_1.kubernetes_version, "1.21.2")
+        # Check if orchestrator_version for each agent pool is updated
+        for agent_pool in updated_mc_1.agent_pool_profiles:
+            self.assertEqual(agent_pool.orchestrator_version, "1.21.2")
+
+        # Second test case with both current_kubernetes_version and current_orchestrator_version as None
+        dec_2 = AKSManagedClusterUpdateDecorator(
+            self.cmd,
+            self.client,
+            raw_param_dict,
+            ResourceType.MGMT_CONTAINERSERVICE,
+        )
+        mc_2 = self.models.ManagedCluster(
+            location="test_location",
+            kubernetes_version="1.18.14",
+            agent_pool_profiles=[
+                self.models.ManagedClusterAgentPoolProfile(
+                    name="agentpool1",
+                    orchestrator_version="1.18.14"
+                ),
+                self.models.ManagedClusterAgentPoolProfile(
+                    name="agentpool2",
+                    orchestrator_version="1.19.7"
+                )
+            ]
+        )
+        # Set current kubernetes version and orchestrator version to None
+        mc_2.current_kubernetes_version = None
+        mc_2.agent_pool_profiles[0].current_orchestrator_version = None
+        mc_2.agent_pool_profiles[1].current_orchestrator_version = None
+
+        dec_2.context.attach_mc(mc_2)
+        with patch(
+            "azure.cli.command_modules.acs.managed_cluster_decorator.prompt_y_n",
+            return_value=True,
+        ):
+            updated_mc_2 = dec_2.update_kubernetes_version_and_orchestrator_version(mc_2)
+
+        # Check if kubernetes_version and orchestrator_version remain unchanged
+        self.assertEqual(updated_mc_2.kubernetes_version, "1.18.14")
+        self.assertEqual(updated_mc_2.agent_pool_profiles[0].orchestrator_version, "1.18.14")
+        self.assertEqual(updated_mc_2.agent_pool_profiles[1].orchestrator_version, "1.19.7")
+
+        # Third test case with both current_kubernetes_version and current_orchestrator_version as None
+        dec_3 = AKSManagedClusterUpdateDecorator(
+            self.cmd,
+            self.client,
+            {},
+            ResourceType.MGMT_CONTAINERSERVICE,
+        )
+        mc_3 = self.models.ManagedCluster(
+            location="test_location",
+            kubernetes_version="1.18.14",
+            agent_pool_profiles=[
+                self.models.ManagedClusterAgentPoolProfile(
+                    name="agentpool1",
+                    orchestrator_version="1.18.14"
+                ),
+                self.models.ManagedClusterAgentPoolProfile(
+                    name="agentpool2",
+                    orchestrator_version="1.19.7"
+                )
+            ]
+        )
+        # Set current kubernetes version and orchestrator version to 1.21.2
+        mc_3.current_kubernetes_version = "1.21.2"
+        mc_3.agent_pool_profiles[0].current_orchestrator_version = "1.21.2"
+        mc_3.agent_pool_profiles[1].current_orchestrator_version = "1.21.2"
+
+        dec_3.context.attach_mc(mc_3) # no prompt 
+        updated_mc_3 = dec_3.update_kubernetes_version_and_orchestrator_version(mc_3)
+
+        # Check if kubernetes_version and orchestrator_version remain unchanged
+        self.assertEqual(updated_mc_3.kubernetes_version, "1.18.14")
+        self.assertEqual(updated_mc_3.agent_pool_profiles[0].orchestrator_version, "1.18.14")
+        self.assertEqual(updated_mc_3.agent_pool_profiles[1].orchestrator_version, "1.19.7")
+
     def test_check_is_postprocessing_required(self):
         dec_1 = AKSManagedClusterUpdateDecorator(
             self.cmd,