From cc0b3d415f43833d212695bf0760b1743c727084 Mon Sep 17 00:00:00 2001
From: Shaoru Hu <49520642+HuShaoRu@users.noreply.github.com>
Date: Thu, 6 Jun 2024 17:31:38 +0800
Subject: [PATCH] {AKS} Set default vm-set-type and ignore spaces for Mixed SKU
 (#7681)

---
 src/aks-preview/HISTORY.rst                       |  1 +
 .../azext_aks_preview/agentpool_decorator.py      | 14 ++++++++++++--
 src/aks-preview/azext_aks_preview/custom.py       | 15 ++++++++++-----
 3 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/src/aks-preview/HISTORY.rst b/src/aks-preview/HISTORY.rst
index 6e3412914b7..a5af0d76d59 100644
--- a/src/aks-preview/HISTORY.rst
+++ b/src/aks-preview/HISTORY.rst
@@ -13,6 +13,7 @@ Pending
 +++++++
 * Update --enable-advanced-network-observability description to note additional costs and add missing flag to create command.
 * Add etag support (--if-match, --if-none-match) to some aks commands for optimistic concurrency control.
+* Change default value of `--vm-set-type` to VirtualMachines when `--vm-sizes` is set
 
 4.0.0b5
 ++++++++
diff --git a/src/aks-preview/azext_aks_preview/agentpool_decorator.py b/src/aks-preview/azext_aks_preview/agentpool_decorator.py
index e1f47acd174..ae73a23b395 100644
--- a/src/aks-preview/azext_aks_preview/agentpool_decorator.py
+++ b/src/aks-preview/azext_aks_preview/agentpool_decorator.py
@@ -86,7 +86,17 @@ def get_vm_set_type(self) -> str:
         :return: string
         """
         # read the original value passed by the command
-        vm_set_type = self.raw_param.get("vm_set_type", CONST_VIRTUAL_MACHINE_SCALE_SETS)
+        vm_set_type = self.raw_param.get("vm_set_type")
+        if vm_set_type is None:
+            if self.raw_param.get("vm_sizes") is None:
+                vm_set_type = CONST_VIRTUAL_MACHINE_SCALE_SETS
+            else:
+                vm_set_type = CONST_VIRTUAL_MACHINES
+        else:
+            if vm_set_type.lower() != CONST_VIRTUAL_MACHINES.lower() and self.raw_param.get("vm_sizes") is not None:
+                raise InvalidArgumentValueError(
+                    "--vm-sizes can only be used with --vm-set-type VirtualMachines(Preview)"
+                )
         # try to read the property value corresponding to the parameter from the `agentpool` object
         if self.agentpool_decorator_mode == AgentPoolDecoratorMode.MANAGED_CLUSTER:
             if self.agentpool and self.agentpool.type is not None:
@@ -640,7 +650,7 @@ def get_vm_sizes(self) -> List[str]:
         """
         raw_value = self.raw_param.get("vm_sizes")
         if raw_value is not None:
-            vm_sizes = raw_value.split(",")
+            vm_sizes = [x.strip() for x in raw_value.split(",")]
         else:
             vm_sizes = [self.get_node_vm_size()]
         return vm_sizes
diff --git a/src/aks-preview/azext_aks_preview/custom.py b/src/aks-preview/azext_aks_preview/custom.py
index 42060b57fb6..f98cd51660b 100644
--- a/src/aks-preview/azext_aks_preview/custom.py
+++ b/src/aks-preview/azext_aks_preview/custom.py
@@ -1772,7 +1772,8 @@ def aks_agentpool_manual_scale_add(cmd,
         resource_type=CUSTOM_MGMT_AKS_PREVIEW,
         operation_group="managed_clusters",
     )
-    new_manual_scale_profile = ManualScaleProfile(sizes=vm_sizes.split(","), count=int(node_count))
+    sizes = [x.strip() for x in vm_sizes.split(",")]
+    new_manual_scale_profile = ManualScaleProfile(sizes=sizes, count=int(node_count))
     instance.virtual_machines_profile.scale.manual.append(new_manual_scale_profile)
 
     return sdk_no_wait(
@@ -1800,12 +1801,15 @@ def aks_agentpool_manual_scale_update(cmd,    # pylint: disable=unused-argument
     instance = client.get(resource_group_name, cluster_name, nodepool_name)
     if instance.type_properties_type != CONST_VIRTUAL_MACHINES:
         raise ClientRequestError("Cannot update manual in a non-virtualmachines node pool.")
+
+    _current_vm_sizes = [x.strip() for x in current_vm_sizes.split(",")]
+    _vm_sizes = [x.strip() for x in vm_sizes.split(",")] if vm_sizes else []
     manual_exists = False
     for m in instance.virtual_machines_profile.scale.manual:
-        if m.sizes == current_vm_sizes.split(","):
+        if m.sizes == _current_vm_sizes:
             manual_exists = True
             if vm_sizes:
-                m.sizes = vm_sizes.split(",")
+                m.sizes = _vm_sizes
             if node_count:
                 m.count = int(node_count)
             break
@@ -1834,15 +1838,16 @@ def aks_agentpool_manual_scale_delete(cmd,    # pylint: disable=unused-argument
     instance = client.get(resource_group_name, cluster_name, nodepool_name)
     if instance.type_properties_type != CONST_VIRTUAL_MACHINES:
         raise CLIError("Cannot delete manual in a non-virtualmachines node pool.")
+    _current_vm_sizes = [x.strip() for x in current_vm_sizes.split(",")]
     manual_exists = False
     for m in instance.virtual_machines_profile.scale.manual:
-        if m.sizes == current_vm_sizes.split(","):
+        if m.sizes == _current_vm_sizes:
             manual_exists = True
             instance.virtual_machines_profile.scale.manual.remove(m)
             break
     if not manual_exists:
         raise InvalidArgumentValueError(
-            f"Manual with sizes {','.join(current_vm_sizes)} doesn't exist in node pool {nodepool_name}"
+            f"Manual with sizes {current_vm_sizes} doesn't exist in node pool {nodepool_name}"
         )
 
     return sdk_no_wait(