From b3a7e26d3867b76b1188d6e548e5c6c2d9ec8e6e Mon Sep 17 00:00:00 2001
From: Fuming Zhang <fumingzhang@microsoft.com>
Date: Wed, 27 Oct 2021 15:37:21 +0800
Subject: [PATCH] add support for node_zones & enable_vmss

---
 .../azext_aks_preview/decorator.py            | 39 ++++++++++-
 .../tests/latest/test_decorator.py            | 70 +++++++++++++++++++
 2 files changed, 108 insertions(+), 1 deletion(-)

diff --git a/src/aks-preview/azext_aks_preview/decorator.py b/src/aks-preview/azext_aks_preview/decorator.py
index af587dea3b5..1876ffd3b1c 100644
--- a/src/aks-preview/azext_aks_preview/decorator.py
+++ b/src/aks-preview/azext_aks_preview/decorator.py
@@ -5,7 +5,7 @@
 
 import os
 import time
-from typing import Dict, Tuple, TypeVar, Union
+from typing import Dict, List, Tuple, TypeVar, Union
 
 from azure.cli.command_modules.acs._consts import (
     DecoratorEarlyExitException,
@@ -121,6 +121,43 @@ def __init__(
     ):
         super().__init__(cmd, raw_parameters, models, decorator_mode)
 
+    # pylint: disable=unused-argument
+    def _get_vm_set_type(self, read_only: bool = False, **kwargs) -> Union[str, None]:
+        """Internal function to dynamically obtain the value of vm_set_type according to the context.
+
+        Note: Inherited and extended in aks-preview to add support for the deprecated option --enable-vmss.
+
+        :return: string or None
+        """
+        vm_set_type = super()._get_vm_set_type(read_only, **kwargs)
+
+        # TODO: Remove the below section when we deprecate the --enable-vmss flag, kept for back-compatibility only.
+        # read the original value passed by the command
+        enable_vmss = self.raw_param.get("enable_vmss")
+
+        if enable_vmss:
+            if vm_set_type and vm_set_type.lower() != "VirtualMachineScaleSets".lower():
+                raise InvalidArgumentValueError(
+                    "--enable-vmss and provided --vm-set-type ({}) are conflicting with each other".format(
+                        vm_set_type
+                    )
+                )
+            vm_set_type = "VirtualMachineScaleSets"
+        return vm_set_type
+
+    def get_zones(self) -> Union[List[str], None]:
+        """Obtain the value of zones.
+
+        Note: Inherited and extended in aks-preview to add support for a different parameter name (node_zones).
+
+        :return: list of strings or None
+        """
+        zones = super().get_zones()
+        if zones is not None:
+            return zones
+        # read the original value passed by the command
+        return self.raw_param.get("node_zones")
+
     def get_pod_subnet_id(self) -> Union[str, None]:
         """Obtain the value of pod_subnet_id.
 
diff --git a/src/aks-preview/azext_aks_preview/tests/latest/test_decorator.py b/src/aks-preview/azext_aks_preview/tests/latest/test_decorator.py
index 154d2d41c42..5d5819cbffb 100644
--- a/src/aks-preview/azext_aks_preview/tests/latest/test_decorator.py
+++ b/src/aks-preview/azext_aks_preview/tests/latest/test_decorator.py
@@ -99,6 +99,76 @@ def setUp(self):
         self.cmd = MockCmd(self.cli_ctx)
         self.models = AKSPreviewModels(self.cmd, CUSTOM_MGMT_AKS_PREVIEW)
 
+    def test__get_vm_set_type(self):
+        # default & dynamic completion
+        ctx_1 = AKSPreviewContext(
+            self.cmd,
+            {"vm_set_type": None, "kubernetes_version": "", "enable_vmss": False},
+            self.models,
+            decorator_mode=DecoratorMode.CREATE,
+        )
+        self.assertEqual(ctx_1._get_vm_set_type(read_only=True), None)
+        self.assertEqual(ctx_1.get_vm_set_type(), "VirtualMachineScaleSets")
+        agent_pool_profile = self.models.ManagedClusterAgentPoolProfile(
+            name="test_ap_name", type="test_mc_vm_set_type"
+        )
+        mc = self.models.ManagedCluster(
+            location="test_location", agent_pool_profiles=[agent_pool_profile]
+        )
+        ctx_1.attach_mc(mc)
+        self.assertEqual(ctx_1.get_vm_set_type(), "test_mc_vm_set_type")
+
+        # custom value & dynamic completion
+        ctx_2 = AKSPreviewContext(
+            self.cmd,
+            {"vm_set_type": "availabilityset", "kubernetes_version": "", "enable_vmss": True},
+            self.models,
+            decorator_mode=DecoratorMode.CREATE,
+        )
+        # fail on invalid vm_set_type when enable_vmss is specified
+        with self.assertRaises(InvalidArgumentValueError):
+            self.assertEqual(ctx_2.get_vm_set_type(), "AvailabilitySet")
+
+        # custom value & dynamic completion
+        ctx_3 = AKSPreviewContext(
+            self.cmd,
+            {"vm_set_type": None, "kubernetes_version": "", "enable_vmss": True},
+            self.models,
+            decorator_mode=DecoratorMode.CREATE,
+        )
+        # fail on invalid vm_set_type when enable_vmss is specified
+        self.assertEqual(ctx_3.get_vm_set_type(), "VirtualMachineScaleSets")
+
+    def test_get_zones(self):
+        # default
+        ctx_1 = AKSPreviewContext(
+            self.cmd,
+            {"node_zones": None},
+            self.models,
+            decorator_mode=DecoratorMode.CREATE,
+        )
+        self.assertEqual(ctx_1.get_zones(), None)
+        agent_pool_profile = self.models.ManagedClusterAgentPoolProfile(
+            name="test_nodepool_name",
+            availability_zones=["test_mc_zones1", "test_mc_zones2"],
+        )
+        mc = self.models.ManagedCluster(
+            location="test_location", agent_pool_profiles=[agent_pool_profile]
+        )
+        ctx_1.attach_mc(mc)
+        self.assertEqual(
+            ctx_1.get_zones(), ["test_mc_zones1", "test_mc_zones2"]
+        )
+
+        # custom value
+        ctx_2 = AKSPreviewContext(
+            self.cmd,
+            {"node_zones": ["test_zones1", "test_zones2"]},
+            self.models,
+            decorator_mode=DecoratorMode.CREATE,
+        )
+        self.assertEqual(ctx_2.get_zones(), ["test_zones1", "test_zones2"])
+
     def test_get_pod_subnet_id(self):
         # default
         ctx_1 = AKSPreviewContext(