update

Author: MuhammadAliFleet

src/fleet/azext_fleet/_client_factory.py

@@ -44,3 +44,8 @@ def cf_auto_upgrade_profiles(cli_ctx, *_):
 
 def cf_auto_upgrade_profile_operations(cli_ctx, *_):
     return get_container_service_client(cli_ctx).auto_upgrade_profile_operations
+
+
+def get_provider_client(cli_ctx):
+    return get_mgmt_service_client(
+        cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES)

src/fleet/azext_fleet/_helpers.py

@@ -13,6 +13,10 @@
 from knack.log import get_logger
 from knack.prompting import NoTTYException, prompt_y_n
 from knack.util import CLIError
+from azure.cli.command_modules.acs._roleassignments import add_role_assignment
+
+from azext_fleet.constants import FLEET_1P_APP_ID
+from azext_fleet._client_factory import get_provider_client
 
 logger = get_logger(__name__)
 
@@ -148,3 +152,15 @@ def _load_kubernetes_configuration(filename):
         raise
     except (yaml.parser.ParserError, UnicodeDecodeError) as ex:
         raise CLIError(f'Error parsing {filename} ({str(ex)})') from ex
+
+
+def validate_subnet(cmd, subnet_id):
+    resource_client = get_provider_client(cmd.cli_ctx)
+    provider = resource_client.providers.get("Microsoft.ContainerService")
+
+    if provider.registration_state != 'Registered':
+        raise CLIError("The Microsoft.ContainerService resource provider is not registered."
+                       "Run `az provider register -n Microsoft.ContainerService --wait`.")
+    if not add_role_assignment(cmd, 'Network Contributor', FLEET_1P_APP_ID, scope=subnet_id):
+        raise CLIError("failed to create role assignment for Fleet RP.\n"
+                      f"Do you have owner permissions on the subnet {subnet_id}?\n")

src/fleet/azext_fleet/custom.py

@@ -10,15 +10,13 @@
 
 from azure.cli.core.commands.client_factory import get_subscription_id
 from azure.cli.core.util import sdk_no_wait
-from azure.cli.command_modules.acs._roleassignments import add_role_assignment
 
 from azext_fleet._client_factory import CUSTOM_MGMT_FLEET
-from azext_fleet._helpers import print_or_merge_credentials
+from azext_fleet._helpers import print_or_merge_credentials, validate_subnet
 from azext_fleet.constants import UPGRADE_TYPE_CONTROLPLANEONLY
 from azext_fleet.constants import UPGRADE_TYPE_FULL
 from azext_fleet.constants import UPGRADE_TYPE_NODEIMAGEONLY
 from azext_fleet.constants import UPGRADE_TYPE_ERROR_MESSAGES
-from azext_fleet.constants import FLEET_1P_APP_ID
 
 
 # pylint: disable=too-many-locals
@@ -112,19 +110,9 @@ def create_fleet(cmd,
     )
 
     if enable_private_cluster:
-        if not add_role_assignment(cmd, 'Network Contributor', FLEET_1P_APP_ID, scope=agent_subnet_id):
-            raise CLIError(
-                "failed to create role assignment for Fleet RP.\nDo you have owner permissions on the vnet?\n"
-                "Is the Microsoft.Containerservice namespace registered?\n"
-                "Please try registering Fleet namespace before retrying: az provider register --namespace 'Microsoft.Containerservice'"  # pylint: disable=line-too-long
-            )
+        validate_subnet(cmd, resource_group_name, agent_subnet_id)
     if enable_vnet_integration:
-        if not add_role_assignment(cmd, 'Network Contributor', FLEET_1P_APP_ID, scope=apiserver_subnet_id):
-            raise CLIError(
-                "failed to create role assignment for Fleet RP.\nDo you have owner permissions on the vnet?\n"
-                "Is the Microsoft.Containerservice namespace registered?\n"
-                "Please try registering Fleet namespace before retrying: az provider register --namespace 'Microsoft.Containerservice'"  # pylint: disable=line-too-long
-            )
+        validate_subnet(cmd, resource_group_name, apiserver_subnet_id)
 
     return sdk_no_wait(no_wait,
                        client.begin_create_or_update,