{AKS} Refactor 3rd part of the create sub-command in aks-preview module

src/aks-preview/azext_aks_preview/_natgateway.py

@@ -3,37 +3,36 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-from .vendored_sdks.azure_mgmt_preview_aks.v2021_08_01.models import ManagedClusterNATGatewayProfile
-from .vendored_sdks.azure_mgmt_preview_aks.v2021_08_01.models import ManagedClusterManagedOutboundIPProfile
 
-
-def create_nat_gateway_profile(managed_outbound_ip_count, idle_timeout):
+def create_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, models):
     """parse and build NAT gateway profile"""
     if not is_nat_gateway_profile_provided(managed_outbound_ip_count, idle_timeout):
         return None
 
-    profile = ManagedClusterNATGatewayProfile()
-    return configure_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, profile)
+    # profile = ManagedClusterNATGatewayProfile()
+    profile = models.get("ManagedClusterNATGatewayProfile")()
+    return configure_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, profile, models)
 
 
-def update_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, profile):
+def update_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, profile, models):
     """parse and update an existing NAT gateway profile"""
     if not is_nat_gateway_profile_provided(managed_outbound_ip_count, idle_timeout):
         return profile
 
-    return configure_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, profile)
+    return configure_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, profile, models)
 
 
 def is_nat_gateway_profile_provided(managed_outbound_ip_count, idle_timeout):
     return any([managed_outbound_ip_count, idle_timeout])
 
 
-def configure_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, profile):
+def configure_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, profile, models):
     """configure a NAT Gateway with customer supplied values"""
     if not profile:
         return profile
 
     if managed_outbound_ip_count:
+        ManagedClusterManagedOutboundIPProfile = models.get("ManagedClusterManagedOutboundIPProfile")
         profile.managed_outbound_ip_profile = ManagedClusterManagedOutboundIPProfile(
             count=managed_outbound_ip_count
         )

src/aks-preview/azext_aks_preview/custom.py

@@ -54,6 +54,7 @@
                                     KeyCredential,
                                     ServicePrincipalCreateParameters,
                                     GetObjectsParameters)
+from azext_aks_preview._client_factory import CUSTOM_MGMT_AKS_PREVIEW
 from .vendored_sdks.azure_mgmt_preview_aks.v2021_08_01.models import (ContainerServiceLinuxProfile,
                                                                       ManagedClusterWindowsProfile,
                                                                       ContainerServiceNetworkProfile,
@@ -987,9 +988,14 @@ def aks_create(cmd,     # pylint: disable=too-many-locals,too-many-statements,to
         load_balancer_outbound_ports,
         load_balancer_idle_timeout)
 
+    from azext_aks_preview.decorator import AKSPreviewModels
+    # store all the models used by nat gateway
+    nat_gateway_models = AKSPreviewModels(cmd, CUSTOM_MGMT_AKS_PREVIEW).nat_gateway_models
     nat_gateway_profile = create_nat_gateway_profile(
         nat_gateway_managed_outbound_ip_count,
-        nat_gateway_idle_timeout)
+        nat_gateway_idle_timeout,
+        models=nat_gateway_models,
+    )
 
     outbound_type = _set_outbound_type(
         outbound_type, vnet_subnet_id, load_balancer_sku, load_balancer_profile)
@@ -1479,10 +1485,15 @@ def aks_update(cmd,     # pylint: disable=too-many-statements,too-many-branches,
             instance.network_profile.load_balancer_profile)
 
     if update_natgw_profile:
+        from azext_aks_preview.decorator import AKSPreviewModels
+        # store all the models used by nat gateway
+        nat_gateway_models = AKSPreviewModels(cmd, CUSTOM_MGMT_AKS_PREVIEW).nat_gateway_models
         instance.network_profile.nat_gateway_profile = update_nat_gateway_profile(
             nat_gateway_managed_outbound_ip_count,
             nat_gateway_idle_timeout,
-            instance.network_profile.nat_gateway_profile)
+            instance.network_profile.nat_gateway_profile,
+            models=nat_gateway_models,
+        )
 
     if attach_acr and detach_acr:
         raise CLIError(

src/aks-preview/azext_aks_preview/decorator.py

@@ -24,6 +24,8 @@
 from azure.cli.core.util import get_file_json
 from knack.log import get_logger
 
+from azext_aks_preview._natgateway import create_nat_gateway_profile
+
 logger = get_logger(__name__)
 
 # type variables
@@ -57,6 +59,31 @@ def __init__(self, cmd: AzCommandsLoader, resource_type: ResourceType):
             resource_type=self.resource_type,
             operation_group="managed_clusters",
         )
+        # init nat gateway models
+        self.init_nat_gateway_models()
+
+    def init_nat_gateway_models(self) -> None:
+        """Initialize models used by nat gateway.
+
+        The models are stored in a dictionary, the key is the model name and the value is the model type.
+
+        :return: None
+        """
+        nat_gateway_models = {}
+        nat_gateway_models["ManagedClusterNATGatewayProfile"] = self.__cmd.get_models(
+            "ManagedClusterNATGatewayProfile",
+            resource_type=self.resource_type,
+            operation_group="managed_clusters",
+        )
+        nat_gateway_models["ManagedClusterManagedOutboundIPProfile"] = self.__cmd.get_models(
+            "ManagedClusterManagedOutboundIPProfile",
+            resource_type=self.resource_type,
+            operation_group="managed_clusters",
+        )
+        self.nat_gateway_models = nat_gateway_models
+        # Note: Uncomment the followings to add these models as class attributes.
+        # for model_name, model_type in nat_gateway_models.items():
+        #     setattr(self, model_name, model_type)
 
 
 # pylint: disable=too-many-public-methods
@@ -286,6 +313,51 @@ def get_node_resource_group(self) -> Union[str, None]:
         # this parameter does not need validation
         return node_resource_group
 
+    def get_nat_gateway_managed_outbound_ip_count(self) -> Union[int, None]:
+        """Obtain the value of nat_gateway_managed_outbound_ip_count.
+
+        :return: string or None
+        """
+        # read the original value passed by the command
+        nat_gateway_managed_outbound_ip_count = self.raw_param.get("nat_gateway_managed_outbound_ip_count")
+        # try to read the property value corresponding to the parameter from the `mc` object
+        if (
+            self.mc and
+            self.mc.network_profile and
+            self.mc.network_profile.nat_gateway_profile and
+            self.mc.network_profile.nat_gateway_profile.managed_outbound_ip_profile and
+            self.mc.network_profile.nat_gateway_profile.managed_outbound_ip_profile.count is not None
+        ):
+            nat_gateway_managed_outbound_ip_count = (
+                self.mc.network_profile.nat_gateway_profile.managed_outbound_ip_profile.count
+            )
+
+        # this parameter does not need dynamic completion
+        # this parameter does not need validation
+        return nat_gateway_managed_outbound_ip_count
+
+    def get_nat_gateway_idle_timeout(self) -> Union[int, None]:
+        """Obtain the value of nat_gateway_idle_timeout.
+
+        :return: string or None
+        """
+        # read the original value passed by the command
+        nat_gateway_idle_timeout = self.raw_param.get("nat_gateway_idle_timeout")
+        # try to read the property value corresponding to the parameter from the `mc` object
+        if (
+            self.mc and
+            self.mc.network_profile and
+            self.mc.network_profile.nat_gateway_profile and
+            self.mc.network_profile.nat_gateway_profile.idle_timeout_in_minutes is not None
+        ):
+            nat_gateway_idle_timeout = (
+                self.mc.network_profile.nat_gateway_profile.idle_timeout_in_minutes
+            )
+
+        # this parameter does not need dynamic completion
+        # this parameter does not need validation
+        return nat_gateway_idle_timeout
+
 
 class AKSPreviewCreateDecorator(AKSCreateDecorator):
     # pylint: disable=super-init-not-called
@@ -318,6 +390,9 @@ def __init__(
     def set_up_agent_pool_profiles(self, mc: ManagedCluster) -> ManagedCluster:
         """Set up agent pool profiles for the ManagedCluster object.
 
+        Call the method of the same name in the parent class to set up agent_pool_profiles, and then set some additional
+        properties on this basis.
+
         :return: the ManagedCluster object
         """
         mc = super().set_up_agent_pool_profiles(mc)
@@ -362,6 +437,30 @@ def set_up_node_resource_group(self, mc: ManagedCluster) -> ManagedCluster:
         mc.node_resource_group = self.context.get_node_resource_group()
         return mc
 
+    def set_up_network_profile(self, mc: ManagedCluster) -> ManagedCluster:
+        """Set up network profile for the ManagedCluster object.
+
+        Call the method of the same name in the parent class to set up network_profile, and then set the
+        nat_gateway_profile on this basis.
+
+        :return: the ManagedCluster object
+        """
+        mc = super().set_up_network_profile(mc)
+        network_profile = mc.network_profile
+
+        # build nat gateway profile, which is part of the network profile
+        nat_gateway_profile = create_nat_gateway_profile(
+            self.context.get_nat_gateway_managed_outbound_ip_count(),
+            self.context.get_nat_gateway_idle_timeout(),
+            models=self.models.nat_gateway_models,
+        )
+
+        load_balancer_sku = self.context.get_load_balancer_sku()
+        if load_balancer_sku != "basic":
+            network_profile.nat_gateway_profile = nat_gateway_profile
+        mc.network_profile = network_profile
+        return mc
+
     def construct_preview_mc_profile(self) -> ManagedCluster:
         """The overall controller used to construct the preview ManagedCluster profile.
 

src/aks-preview/azext_aks_preview/tests/latest/test_decorator.py

@@ -294,6 +294,58 @@ def test_get_node_resource_group(self):
             ctx_1.get_node_resource_group(), "test_node_resource_group"
         )
 
+    def test_get_nat_gateway_managed_outbound_ip_count(self):
+        # default
+        ctx_1 = AKSPreviewContext(
+            self.cmd,
+            {"nat_gateway_managed_outbound_ip_count": None},
+            self.models,
+            decorator_mode=DecoratorMode.CREATE,
+        )
+        self.assertEqual(
+            ctx_1.get_nat_gateway_managed_outbound_ip_count(), None
+        )
+        nat_gateway_profile = self.models.nat_gateway_models.get(
+            "ManagedClusterNATGatewayProfile"
+        )(
+            managed_outbound_ip_profile=self.models.nat_gateway_models.get(
+                "ManagedClusterManagedOutboundIPProfile"
+            )(count=10)
+        )
+        network_profile = self.models.ContainerServiceNetworkProfile(
+            nat_gateway_profile=nat_gateway_profile
+        )
+        mc = self.models.ManagedCluster(
+            location="test_location",
+            network_profile=network_profile,
+        )
+        ctx_1.attach_mc(mc)
+        self.assertEqual(ctx_1.get_nat_gateway_managed_outbound_ip_count(), 10)
+
+    def test_get_nat_gateway_idle_timeout(self):
+        # default
+        ctx_1 = AKSPreviewContext(
+            self.cmd,
+            {"nat_gateway_idle_timeout": None},
+            self.models,
+            decorator_mode=DecoratorMode.CREATE,
+        )
+        self.assertEqual(ctx_1.get_nat_gateway_idle_timeout(), None)
+        nat_gateway_profile = self.models.nat_gateway_models.get(
+            "ManagedClusterNATGatewayProfile"
+        )(
+            idle_timeout_in_minutes=20,
+        )
+        network_profile = self.models.ContainerServiceNetworkProfile(
+            nat_gateway_profile=nat_gateway_profile
+        )
+        mc = self.models.ManagedCluster(
+            location="test_location",
+            network_profile=network_profile,
+        )
+        ctx_1.attach_mc(mc)
+        self.assertEqual(ctx_1.get_nat_gateway_idle_timeout(), 20)
+
 
 class AKSPreviewCreateDecoratorTestCase(unittest.TestCase):
     def setUp(self):
@@ -537,6 +589,98 @@ def test_set_up_node_resource_group(self):
         )
         self.assertEqual(dec_mc_2, ground_truth_mc_2)
 
+    def test_set_up_network_profile(self):
+        # default value in `aks_create`
+        dec_1 = AKSPreviewCreateDecorator(
+            self.cmd,
+            self.client,
+            {
+                "load_balancer_sku": None,
+                "load_balancer_managed_outbound_ip_count": None,
+                "load_balancer_outbound_ips": None,
+                "load_balancer_outbound_ip_prefixes": None,
+                "load_balancer_outbound_ports": None,
+                "load_balancer_idle_timeout": None,
+                "outbound_type": None,
+                "network_plugin": None,
+                "pod_cidr": None,
+                "service_cidr": None,
+                "dns_service_ip": None,
+                "docker_bridge_cidr": None,
+                "network_policy": None,
+                "nat_gateway_managed_outbound_ip_count": None,
+                "nat_gateway_idle_timeout": None,
+            },
+            CUSTOM_MGMT_AKS_PREVIEW,
+        )
+
+        mc_1 = self.models.ManagedCluster(location="test_location")
+        # fail on passing the wrong mc object
+        with self.assertRaises(CLIInternalError):
+            dec_1.set_up_network_profile(None)
+        dec_mc_1 = dec_1.set_up_network_profile(mc_1)
+
+        network_profile_1 = self.models.ContainerServiceNetworkProfile(
+            network_plugin="kubenet",  # default value in SDK
+            pod_cidr="10.244.0.0/16",  # default value in SDK
+            service_cidr="10.0.0.0/16",  # default value in SDK
+            dns_service_ip="10.0.0.10",  # default value in SDK
+            docker_bridge_cidr="172.17.0.1/16",  # default value in SDK
+            load_balancer_sku="standard",
+            outbound_type="loadBalancer",
+        )
+        ground_truth_mc_1 = self.models.ManagedCluster(
+            location="test_location", network_profile=network_profile_1
+        )
+        self.assertEqual(dec_mc_1, ground_truth_mc_1)
+
+        # custom value
+        dec_2 = AKSPreviewCreateDecorator(
+            self.cmd,
+            self.client,
+            {
+                "load_balancer_sku": None,
+                "load_balancer_managed_outbound_ip_count": None,
+                "load_balancer_outbound_ips": None,
+                "load_balancer_outbound_ip_prefixes": None,
+                "load_balancer_outbound_ports": None,
+                "load_balancer_idle_timeout": None,
+                "outbound_type": None,
+                "network_plugin": "kubenet",
+                "pod_cidr": "10.246.0.0/16",
+                "service_cidr": None,
+                "dns_service_ip": None,
+                "docker_bridge_cidr": None,
+                "network_policy": None,
+                "nat_gateway_managed_outbound_ip_count": 10,
+                "nat_gateway_idle_timeout": 20,
+            },
+            CUSTOM_MGMT_AKS_PREVIEW,
+        )
+        mc_2 = self.models.ManagedCluster(location="test_location")
+        dec_mc_2 = dec_2.set_up_network_profile(mc_2)
+
+        nat_gateway_profile_2 = self.models.nat_gateway_models.get("ManagedClusterNATGatewayProfile")(
+            managed_outbound_ip_profile=self.models.nat_gateway_models.get(
+                "ManagedClusterManagedOutboundIPProfile"
+            )(count=10),
+            idle_timeout_in_minutes=20,
+        )
+        network_profile_2 = self.models.ContainerServiceNetworkProfile(
+            network_plugin="kubenet",
+            pod_cidr="10.246.0.0/16",
+            service_cidr=None,  # overwritten to None
+            dns_service_ip=None,  # overwritten to None
+            docker_bridge_cidr=None,  # overwritten to None
+            load_balancer_sku="standard",
+            outbound_type="loadBalancer",
+            nat_gateway_profile=nat_gateway_profile_2,
+        )
+        ground_truth_mc_2 = self.models.ManagedCluster(
+            location="test_location", network_profile=network_profile_2
+        )
+        self.assertEqual(dec_mc_2, ground_truth_mc_2)
+
     def test_construct_preview_mc_profile(self):
         pass