{AKS} Add v2 agentpool decorator

src/aks-preview/azcli_aks_live_test/configs/ext_matrix_default.json

@@ -24,7 +24,8 @@
             "test_aks_nodepool_add_with_gpu_instance_profile",
             "test_aks_create_with_crg_id",
             "test_aks_create_and_update_with_http_proxy_config",
-            "test_aks_snapshot"
+            "test_aks_snapshot",
+            "test_aks_custom_kubelet_identity"
         ]
     }
 }

src/aks-preview/azext_aks_preview/_client_factory.py

@@ -14,48 +14,46 @@
                                              'ContainerServiceClient')
 
 
-def cf_storage(cli_ctx, subscription_id=None):
-    return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_STORAGE, subscription_id=subscription_id)
-
-
-def cf_compute_service(cli_ctx, *_):
-    return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_COMPUTE)
+def get_container_service_client(cli_ctx, **_):
+    return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW)
 
 
 def cf_container_services(cli_ctx, *_):
     return get_container_service_client(cli_ctx).container_services
 
 
-def get_container_service_client(cli_ctx, **_):
-    return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW)
-
-
 def cf_managed_clusters(cli_ctx, *_):
-    return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW).managed_clusters
+    return get_container_service_client(cli_ctx).managed_clusters
 
 
 def cf_agent_pools(cli_ctx, *_):
-    return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW).agent_pools
+    return get_container_service_client(cli_ctx).agent_pools
 
 
-def cf_mc_snapshots_client(cli_ctx, subscription_id=None):
-    return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW, subscription_id=subscription_id).managed_cluster_snapshots
+def cf_maintenance_configurations(cli_ctx, *_):
+    return get_container_service_client(cli_ctx).maintenance_configurations
 
 
-def cf_mc_snapshots(cli_ctx, *_):
-    return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW).managed_cluster_snapshots
+def cf_nodepool_snapshots(cli_ctx, *_):
+    return get_container_service_client(cli_ctx).snapshots
 
 
+# TODO: remove this
 def cf_nodepool_snapshots_client(cli_ctx, subscription_id=None):
     return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW, subscription_id=subscription_id).snapshots
 
 
-def cf_nodepool_snapshots(cli_ctx, *_):
-    return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW).snapshots
+def cf_mc_snapshots(cli_ctx, *_):
+    return get_container_service_client(cli_ctx).managed_cluster_snapshots
 
 
-def cf_maintenance_configurations(cli_ctx, *_):
-    return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW).maintenance_configurations
+# TODO: remove this
+def cf_mc_snapshots_client(cli_ctx, subscription_id=None):
+    return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW, subscription_id=subscription_id).managed_cluster_snapshots
+
+
+def cf_compute_service(cli_ctx, *_):
+    return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_COMPUTE)
 
 
 def cf_resource_groups(cli_ctx, subscription_id=None):
@@ -73,6 +71,10 @@ def cf_container_registry_service(cli_ctx, subscription_id=None):
                                    subscription_id=subscription_id)
 
 
+def cf_storage(cli_ctx, subscription_id=None):
+    return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_STORAGE, subscription_id=subscription_id)
+
+
 def get_auth_management_client(cli_ctx, scope=None, **_):
     import re
 

src/aks-preview/azext_aks_preview/_help.py

@@ -30,6 +30,9 @@
           type: bool
           short-summary: Skip role assignment for subnet (advanced networking).
           long-summary:  If specified, please make sure your service principal has the access to your subnet.
+        - name: --zones -z
+          type: string array
+          short-summary: Space-separated list of availability zones where agent nodes will be placed.
         - name: --client-secret
           type: string
           short-summary: Secret associated with the service principal. This argument is required if
@@ -990,9 +993,9 @@
           type: int
           short-summary: The maximum number of pods deployable to a node.
           long-summary: If not specified, defaults based on network-plugin. 30 for "azure", 110 for "kubenet", or 250 for "none".
-        - name: --node-zones --zones -z
+        - name: --zones -z
           type: string array
-          short-summary: (will be deprecated, use --zones) Availability zones where agent nodes will be placed.
+          short-summary: Space-separated list of availability zones where agent nodes will be placed.
         - name: --vnet-subnet-id
           type: string
           short-summary: The ID of a subnet in an existing VNet into which to deploy the cluster.

src/aks-preview/azext_aks_preview/_helpers.py

@@ -2,7 +2,14 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
-from ._consts import CONST_CONTAINER_NAME_MAX_LENGTH
+import re
+
+from azure.cli.command_modules.acs._helpers import map_azure_error_to_cli_error
+from azure.cli.core.azclierror import InvalidArgumentValueError, ResourceNotFoundError
+from azure.core.exceptions import AzureError
+
+from azext_aks_preview._client_factory import cf_nodepool_snapshots
+from azext_aks_preview._consts import CONST_CONTAINER_NAME_MAX_LENGTH
 
 
 def _trim_fqdn_name_containing_hcp(normalized_fqdn: str) -> str:
@@ -65,3 +72,29 @@ def similar_word(a, b):
             matches.append(word)
 
     return matches
+
+
+def get_nodepool_snapshot_by_snapshot_id(cli_ctx, snapshot_id):
+    _re_snapshot_resource_id = re.compile(
+        r"/subscriptions/(.*?)/resourcegroups/(.*?)/providers/microsoft.containerservice/snapshots/(.*)",
+        flags=re.IGNORECASE,
+    )
+    snapshot_id = snapshot_id.lower()
+    match = _re_snapshot_resource_id.search(snapshot_id)
+    if match:
+        resource_group_name = match.group(2)
+        snapshot_name = match.group(3)
+        return get_nodepool_snapshot(cli_ctx, resource_group_name, snapshot_name)
+    raise InvalidArgumentValueError("Cannot parse snapshot name from provided resource id '{}'.".format(snapshot_id))
+
+
+def get_nodepool_snapshot(cli_ctx, resource_group_name, snapshot_name):
+    snapshot_client = cf_nodepool_snapshots(cli_ctx)
+    try:
+        snapshot = snapshot_client.get(resource_group_name, snapshot_name)
+    # track 2 sdk raise exception from azure.core.exceptions
+    except AzureError as ex:
+        if "not found" in ex.message:
+            raise ResourceNotFoundError("Snapshot '{}' not found.".format(snapshot_name))
+        raise map_azure_error_to_cli_error(ex)
+    return snapshot