{AKS} Refactor aks nodepool update - part 1 (#22050)

* sort params

* refactor auto scaler, label, tag, taint

* add test cases for getters

* add more test cases

* add more tests

* fix lint

Author: FumingZhang

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

@@ -448,16 +448,16 @@ def load_arguments(self, _):
         c.argument('enable_cluster_autoscaler', options_list=["--enable-cluster-autoscaler", "-e"], action='store_true')
         c.argument('min_count', type=int, validator=validate_nodes_count)
         c.argument('max_count', type=int, validator=validate_nodes_count)
-        c.argument('scale_down_mode', arg_type=get_enum_type([CONST_SCALE_DOWN_MODE_DELETE, CONST_SCALE_DOWN_MODE_DEALLOCATE]))
         c.argument('priority', arg_type=get_enum_type(node_priorities), validator=validate_priority)
         c.argument('eviction_policy', arg_type=get_enum_type(node_eviction_policies), validator=validate_eviction_policy)
         c.argument('spot_max_price', type=float, validator=validate_spot_max_price)
         c.argument('labels', nargs='*', validator=validate_nodepool_labels)
         c.argument('tags', tags_type)
         c.argument('node_taints', validator=validate_taints)
-        c.argument('mode', get_enum_type(node_mode_types))
         c.argument('node_osdisk_type', arg_type=get_enum_type(node_os_disk_types))
         c.argument('node_osdisk_size', type=int)
+        c.argument('mode', get_enum_type(node_mode_types))
+        c.argument('scale_down_mode', arg_type=get_enum_type([CONST_SCALE_DOWN_MODE_DELETE, CONST_SCALE_DOWN_MODE_DEALLOCATE]))
         c.argument('max_surge', validator=validate_max_surge)
         c.argument('max_pods', type=int, options_list=['--max-pods', '-m'])
         c.argument('zones', zones_type, options_list=['--zones', '-z'], help='Space-separated list of availability zones where agent nodes will be placed.')
@@ -478,11 +478,11 @@ def load_arguments(self, _):
                    "--update-cluster-autoscaler", "-u"], action='store_true')
         c.argument('min_count', type=int, validator=validate_nodes_count)
         c.argument('max_count', type=int, validator=validate_nodes_count)
-        c.argument('scale_down_mode', arg_type=get_enum_type([CONST_SCALE_DOWN_MODE_DELETE, CONST_SCALE_DOWN_MODE_DEALLOCATE]))
         c.argument('labels', nargs='*', validator=validate_nodepool_labels)
         c.argument('tags', tags_type)
         c.argument('node_taints', validator=validate_taints)
         c.argument('mode', get_enum_type(node_mode_types))
+        c.argument('scale_down_mode', arg_type=get_enum_type([CONST_SCALE_DOWN_MODE_DELETE, CONST_SCALE_DOWN_MODE_DEALLOCATE]))
         c.argument('max_surge', validator=validate_max_surge)
 
     with self.argument_context('aks nodepool upgrade') as c:

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

@@ -20,13 +20,21 @@
     CONST_SPOT_EVICTION_POLICY_DELETE,
     CONST_VIRTUAL_MACHINE_SCALE_SETS,
     AgentPoolDecoratorMode,
+    DecoratorEarlyExitException,
     DecoratorMode,
 )
 from azure.cli.command_modules.acs._helpers import get_snapshot_by_snapshot_id
 from azure.cli.command_modules.acs._validators import extract_comma_separated_string
 from azure.cli.command_modules.acs.base_decorator import BaseAKSContext, BaseAKSModels, BaseAKSParamDict
+from azure.cli.command_modules.acs.decorator import safe_list_get
 from azure.cli.core import AzCommandsLoader
-from azure.cli.core.azclierror import CLIInternalError, InvalidArgumentValueError, RequiredArgumentMissingError
+from azure.cli.core.azclierror import (
+    ArgumentUsageError,
+    CLIInternalError,
+    InvalidArgumentValueError,
+    MutuallyExclusiveArgumentError,
+    RequiredArgumentMissingError,
+)
 from azure.cli.core.commands import AzCliCommand
 from azure.cli.core.profiles import ResourceType
 from azure.cli.core.util import get_file_json, sdk_no_wait
@@ -93,6 +101,7 @@ def __init__(
         super().__init__(cmd, raw_parameters, models, decorator_mode)
         self.agentpool_decorator_mode = agentpool_decorator_mode
         self.agentpool = None
+        self._agentpools = []
 
     # pylint: disable=no-self-use
     def __validate_counts_in_autoscaler(
@@ -218,7 +227,10 @@ def _get_nodepool_name(self, enable_validation: bool = False) -> str:
         # this parameter does not need dynamic completion
         # validation
         if enable_validation:
-            if self.agentpool_decorator_mode == AgentPoolDecoratorMode.STANDALONE:
+            if (
+                self.agentpool_decorator_mode == AgentPoolDecoratorMode.STANDALONE and
+                self.decorator_mode == DecoratorMode.CREATE
+            ):
                 agentpool_client = cf_agent_pools(self.cmd.cli_ctx)
                 instances = agentpool_client.list(self.get_resource_group_name(), self.get_cluster_name())
                 for agentpool_profile in instances:
@@ -265,7 +277,6 @@ def get_max_surge(self):
         # this parameter does not need validation
         return max_surge
 
-    # pylint: disable=too-many-branches
     def get_node_count_and_enable_cluster_autoscaler_min_max_count(
         self,
     ) -> Tuple[int, bool, Union[int, None], Union[int, None]]:
@@ -316,6 +327,93 @@ def get_node_count_and_enable_cluster_autoscaler_min_max_count(
         )
         return node_count, enable_cluster_autoscaler, min_count, max_count
 
+    def get_update_enable_disable_cluster_autoscaler_and_min_max_count(
+        self,
+    ) -> Tuple[bool, bool, bool, Union[int, None], Union[int, None]]:
+        """Obtain the value of update_cluster_autoscaler, enable_cluster_autoscaler, disable_cluster_autoscaler,
+        min_count and max_count.
+
+        This function will verify the parameters through function "__validate_counts_in_autoscaler" by default. Besides
+        if both enable_cluster_autoscaler and update_cluster_autoscaler are specified, a MutuallyExclusiveArgumentError
+        will be raised. If enable_cluster_autoscaler or update_cluster_autoscaler is specified and there are multiple
+        agent pool profiles, an ArgumentUsageError will be raised. If enable_cluster_autoscaler is specified and
+        autoscaler is already enabled in `mc`, it will output warning messages and exit with code 0. If
+        update_cluster_autoscaler is specified and autoscaler is not enabled in `mc`, it will raise an
+        InvalidArgumentValueError. If disable_cluster_autoscaler is specified and autoscaler is not enabled in `mc`,
+        it will output warning messages and exit with code 0.
+
+        :return: a tuple containing four elements: update_cluster_autoscaler of bool type, enable_cluster_autoscaler
+        of bool type, disable_cluster_autoscaler of bool type, min_count of int type or None and max_count of int type
+        or None
+        """
+        # update_cluster_autoscaler
+        # read the original value passed by the command
+        update_cluster_autoscaler = self.raw_param.get("update_cluster_autoscaler", False)
+
+        # enable_cluster_autoscaler
+        # read the original value passed by the command
+        enable_cluster_autoscaler = self.raw_param.get("enable_cluster_autoscaler", False)
+
+        # disable_cluster_autoscaler
+        # read the original value passed by the command
+        disable_cluster_autoscaler = self.raw_param.get("disable_cluster_autoscaler", False)
+
+        # min_count
+        # read the original value passed by the command
+        min_count = self.raw_param.get("min_count")
+
+        # max_count
+        # read the original value passed by the command
+        max_count = self.raw_param.get("max_count")
+
+        # these parameters do not need dynamic completion
+
+        # validation
+        if self.agentpool_decorator_mode == AgentPoolDecoratorMode.MANAGED_CLUSTER:
+            # For multi-agent pool, use the az aks nodepool command
+            if (enable_cluster_autoscaler or update_cluster_autoscaler) and len(self._agentpools) > 1:
+                raise ArgumentUsageError(
+                    'There are more than one node pool in the cluster. Please use "az aks nodepool" command '
+                    "to update per node pool auto scaler settings"
+                )
+
+        if enable_cluster_autoscaler + update_cluster_autoscaler + disable_cluster_autoscaler > 1:
+            raise MutuallyExclusiveArgumentError(
+                "Can only specify one of --enable-cluster-autoscaler, --update-cluster-autoscaler and "
+                "--disable-cluster-autoscaler"
+            )
+
+        self.__validate_counts_in_autoscaler(
+            None,
+            enable_cluster_autoscaler or update_cluster_autoscaler,
+            min_count,
+            max_count,
+            decorator_mode=DecoratorMode.UPDATE,
+        )
+
+        if enable_cluster_autoscaler and self.agentpool.enable_auto_scaling:
+            logger.warning(
+                "Cluster autoscaler is already enabled for this node pool.\n"
+                'Please run "az aks --update-cluster-autoscaler" '
+                "if you want to update min-count or max-count."
+            )
+            raise DecoratorEarlyExitException()
+
+        if update_cluster_autoscaler and not self.agentpool.enable_auto_scaling:
+            raise InvalidArgumentValueError(
+                "Cluster autoscaler is not enabled for this node pool.\n"
+                'Run "az aks nodepool update --enable-cluster-autoscaler" '
+                "to enable cluster with min-count and max-count."
+            )
+
+        if disable_cluster_autoscaler and not self.agentpool.enable_auto_scaling:
+            logger.warning(
+                "Cluster autoscaler is already disabled for this node pool."
+            )
+            raise DecoratorEarlyExitException()
+
+        return update_cluster_autoscaler, enable_cluster_autoscaler, disable_cluster_autoscaler, min_count, max_count
+
     def get_node_osdisk_size(self) -> Union[int, None]:
         """Obtain the value of node_osdisk_size.
 
@@ -583,9 +681,10 @@ def get_nodepool_labels(self) -> Union[Dict[str, str], None]:
         else:
             nodepool_labels = self.raw_param.get("labels")
 
-        # try to read the property value corresponding to the parameter from the `agentpool` object
-        if self.agentpool and self.agentpool.node_labels is not None:
-            nodepool_labels = self.agentpool.node_labels
+        # In create mode, try to read the property value corresponding to the parameter from the `agentpool` object
+        if self.decorator_mode == DecoratorMode.CREATE:
+            if self.agentpool and self.agentpool.node_labels is not None:
+                nodepool_labels = self.agentpool.node_labels
 
         # this parameter does not need dynamic completion
         # this parameter does not need validation
@@ -602,9 +701,10 @@ def get_nodepool_tags(self) -> Union[Dict[str, str], None]:
         else:
             nodepool_tags = self.raw_param.get("tags")
 
-        # try to read the property value corresponding to the parameter from the `agentpool` object
-        if self.agentpool and self.agentpool.tags is not None:
-            nodepool_tags = self.agentpool.tags
+        # In create mode, try to read the property value corresponding to the parameter from the `agentpool` object
+        if self.decorator_mode == DecoratorMode.CREATE:
+            if self.agentpool and self.agentpool.tags is not None:
+                nodepool_tags = self.agentpool.tags
 
         # this parameter does not need dynamic completion
         # this parameter does not need validation
@@ -617,13 +717,17 @@ def get_node_taints(self) -> Union[List[str], None]:
         """
         # read the original value passed by the command
         node_taints = self.raw_param.get("node_taints")
-        # normalize, keep None as None
+        # normalize, default is an empty list
         if node_taints is not None:
             node_taints = [x.strip() for x in (node_taints.split(",") if node_taints else [])]
+        # keep None as None for update mode
+        if node_taints is None and self.decorator_mode == DecoratorMode.CREATE:
+            node_taints = []
 
-        # try to read the property value corresponding to the parameter from the `agentpool` object
-        if self.agentpool and self.agentpool.node_taints is not None:
-            node_taints = self.agentpool.node_taints
+        # In create mode, try to read the property value corresponding to the parameter from the `agentpool` object
+        if self.decorator_mode == DecoratorMode.CREATE:
+            if self.agentpool and self.agentpool.node_taints is not None:
+                node_taints = self.agentpool.node_taints
 
         # this parameter does not need validation
         return node_taints
@@ -1228,8 +1332,8 @@ def set_up_custom_node_config(self, agentpool: AgentPool) -> AgentPool:
         agentpool.linux_os_config = self.context.get_linux_os_config()
         return agentpool
 
-    def construct_default_agentpool_profile(self) -> AgentPool:
-        """The overall controller used to construct the default AgentPool profile.
+    def construct_agentpool_profile_default(self) -> AgentPool:
+        """The overall controller used to construct the AgentPool profile by default.
 
         The completely constructed AgentPool object will later be passed as a parameter to the underlying SDK
         (mgmt-containerservice) to send the actual request.
@@ -1266,8 +1370,8 @@ def construct_default_agentpool_profile(self) -> AgentPool:
     def add_agentpool(self, agentpool: AgentPool) -> AgentPool:
         """Send request to add a new agentpool.
 
-        The function "sdk_no_wait" will be called to use the ContainerServiceClient to send a reqeust to add a new agent
-        pool to the cluster.
+        The function "sdk_no_wait" will be called to use the Agentpool operations of ContainerServiceClient to send a
+        reqeust to add a new agent pool to the cluster.
 
         :return: the AgentPool object
         """
@@ -1292,6 +1396,7 @@ def __init__(
         client: AgentPoolsOperations,
         raw_parameters: Dict,
         resource_type: ResourceType,
+        agentpool_decorator_mode: AgentPoolDecoratorMode,
     ):
         """Internal controller of aks_agentpool_update.
 
@@ -1303,8 +1408,135 @@ def __init__(
         """
         self.cmd = cmd
         self.client = client
-        self.models = AKSAgentPoolModels(cmd, resource_type)
+        self.agentpool_decorator_mode = agentpool_decorator_mode
+        self.models = AKSAgentPoolModels(cmd, resource_type, agentpool_decorator_mode)
         # store the context in the process of assemble the AgentPool object
         self.context = AKSAgentPoolContext(
-            cmd, AKSAgentPoolParamDict(raw_parameters), self.models, decorator_mode=DecoratorMode.UPDATE
+            cmd, AKSAgentPoolParamDict(raw_parameters), self.models, DecoratorMode.UPDATE, agentpool_decorator_mode
+        )
+
+    def _ensure_agentpool(self, agentpool: AgentPool) -> None:
+        """Internal function to ensure that the incoming `agentpool` object is valid and the same as the attached
+        `agentpool` object in the context.
+
+        If the incoming `agentpool` is not valid or is inconsistent with the `agentpool` in the context, raise a
+        CLIInternalError.
+
+        :return: None
+        """
+        if not isinstance(agentpool, self.models.UnifiedAgentPoolModel):
+            raise CLIInternalError(
+                "Unexpected agentpool object with type '{}'.".format(type(agentpool))
+            )
+
+        if self.context.agentpool != agentpool:
+            raise CLIInternalError(
+                "Inconsistent state detected. The incoming `agentpool` "
+                "is not the same as the `agentpool` in the context."
+            )
+
+    # pylint: disable=protected-access
+    def fetch_agentpool(self, agentpools: List[AgentPool] = None) -> AgentPool:
+        """Get the AgentPool object currently in use and attach it to internal context.
+
+        Internally send request using Agentpool operations of ContainerServiceClient.
+
+        :return: the AgentPool object
+        """
+        if self.agentpool_decorator_mode == AgentPoolDecoratorMode.MANAGED_CLUSTER:
+            self.context._agentpools = agentpools
+            agentpool = safe_list_get(agentpools, 0)
+        else:
+            agentpool = self.client.get(
+                self.context.get_resource_group_name(),
+                self.context.get_cluster_name(),
+                self.context.get_nodepool_name(),
+            )
+
+        # attach agentpool to AKSAgentPoolContext
+        self.context.attach_agentpool(agentpool)
+        return agentpool
+
+    def update_auto_scaler_properties(self, agentpool: AgentPool) -> AgentPool:
+        """Update auto scaler related properties for the Agentpool object.
+
+        :return: the Agentpool object
+        """
+        self._ensure_agentpool(agentpool)
+
+        (
+            update_cluster_autoscaler,
+            enable_cluster_autoscaler,
+            disable_cluster_autoscaler,
+            min_count,
+            max_count,
+        ) = (
+            self.context.get_update_enable_disable_cluster_autoscaler_and_min_max_count()
+        )
+
+        if update_cluster_autoscaler or enable_cluster_autoscaler:
+            agentpool.enable_auto_scaling = True
+            agentpool.min_count = min_count
+            agentpool.max_count = max_count
+
+        if disable_cluster_autoscaler:
+            agentpool.enable_auto_scaling = False
+            agentpool.min_count = None
+            agentpool.max_count = None
+        return agentpool
+
+    def update_label_tag_taint(self, agentpool: AgentPool) -> AgentPool:
+        """Set up label, tag, taint for the AgentPool object.
+
+        :return: the AgentPool object
+        """
+        self._ensure_agentpool(agentpool)
+
+        labels = self.context.get_nodepool_labels()
+        if labels is not None:
+            agentpool.node_labels = labels
+
+        tags = self.context.get_nodepool_tags()
+        if tags is not None:
+            agentpool.tags = tags
+
+        node_taints = self.context.get_node_taints()
+        if node_taints is not None:
+            agentpool.node_taints = node_taints
+        return agentpool
+
+    def update_agentpool_profile_default(self, agentpools: List[AgentPool] = None) -> AgentPool:
+        """The overall controller used to update AgentPool profile by default.
+
+        The completely constructed AgentPool object will later be passed as a parameter to the underlying SDK
+        (mgmt-containerservice) to send the actual request.
+
+        :return: the AgentPool object
+        """
+        # fetch the Agentpool object
+        agentpool = self.fetch_agentpool(agentpools)
+        # update auto scaler properties
+        agentpool = self.update_auto_scaler_properties(agentpool)
+        # update label, tag, taint
+        agentpool = self.update_label_tag_taint(agentpool)
+        return agentpool
+
+    def update_agentpool(self, agentpool: AgentPool) -> AgentPool:
+        """Send request to add a new agentpool.
+
+        The function "sdk_no_wait" will be called to use the Agentpool operations of ContainerServiceClient to send a
+        reqeust to update an existing agent pool of the cluster.
+
+        :return: the AgentPool object
+        """
+        self._ensure_agentpool(agentpool)
+
+        return sdk_no_wait(
+            self.context.get_no_wait(),
+            self.client.begin_create_or_update,
+            self.context.get_resource_group_name(),
+            self.context.get_cluster_name(),
+            self.context.get_nodepool_name(),
+            agentpool,
+            headers=self.context.get_aks_custom_headers(),
         )