address comments

Author: MuhammadAliFleet

src/fleet/azext_fleet/_helpers.py

@@ -14,7 +14,7 @@
 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 azure.mgmt.core.tools import is_valid_resource_id, parse_resource_id
+from azure.mgmt.core.tools import parse_resource_id
 
 
 from azext_fleet._client_factory import get_provider_client
@@ -156,32 +156,24 @@ def _load_kubernetes_configuration(filename):
         raise CLIError(f'Error parsing {filename} ({str(ex)})') from ex
 
 
-def assign_network_contributor_role_to_subnet(cmd, objectId, subnet_id):
-    resource_client = get_provider_client(cmd.cli_ctx)
-    provider = resource_client.providers.get("Microsoft.ContainerService")
-
-    # provider registration state being is checked to ensure that the Fleet service principal is available
-    # to create the role assignment on the subnet
-    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', objectId, scope=subnet_id):
+def assign_network_contributor_role_to_subnet(cmd, object_id, subnet_id):
+    if not add_role_assignment(cmd, 'Network Contributor', object_id, scope=subnet_id):
         logger.warning("Failed to create Network Contributor role assignment on the subnet.\n"
                        "Please ensure you have sufficient permissions to assign roles on subnet %s.", subnet_id)
 
 
 def get_msi_object_id(cmd, msi_resource_id):
-    try:
-        if not is_valid_resource_id(msi_resource_id):
-            raise CLIError(f"The provided managed identity resource ID '{msi_resource_id}' is not valid.")
-        parsed = parse_resource_id(msi_resource_id)
-        subscription_id = parsed['subscription']
-        resource_group_name = parsed['resource_group']
-        msi_name = parsed['resource_name']
-        msi_client = get_msi_client(cmd.cli_ctx, subscription_id=subscription_id)
-        msi = msi_client.user_assigned_identities.get(resource_name=msi_name,
-                                                      resource_group_name=resource_group_name)
-        return msi.principal_id
-    except Exception as ex:
-        raise CLIError(f"Failed to get object ID for managed identity {msi_resource_id}: {str(ex)}") from ex
+    parsed = parse_resource_id(msi_resource_id)
+    subscription_id = parsed['subscription']
+    resource_group_name = parsed['resource_group']
+    msi_name = parsed['resource_name']
+    msi_client = get_msi_client(cmd.cli_ctx, subscription_id=subscription_id)
+    msi = msi_client.user_assigned_identities.get(resource_name=msi_name,
+                                                  resource_group_name=resource_group_name)
+    return msi.principal_id
+
+
+def is_rp_registered(cmd):
+    resource_client = get_provider_client(cmd.cli_ctx)
+    provider = resource_client.providers.get("Microsoft.ContainerService")
+    return provider.registration_state == 'Registered'

src/fleet/azext_fleet/custom.py

@@ -12,7 +12,7 @@
 from azure.cli.core.util import sdk_no_wait
 
 from azext_fleet._client_factory import CUSTOM_MGMT_FLEET
-from azext_fleet._helpers import print_or_merge_credentials
+from azext_fleet._helpers import is_rp_registered, print_or_merge_credentials
 from azext_fleet._helpers import assign_network_contributor_role_to_subnet
 from azext_fleet._helpers import get_msi_object_id
 from azext_fleet.constants import UPGRADE_TYPE_CONTROLPLANEONLY
@@ -91,6 +91,7 @@ def create_fleet(cmd,
         resource_type=CUSTOM_MGMT_FLEET,
         operation_group="fleets"
     )
+
     managed_service_identity = fleet_managed_service_identity_model(type="None")
     if enable_managed_identity:
         managed_service_identity.type = "SystemAssigned"
@@ -105,6 +106,11 @@ def create_fleet(cmd,
     elif assign_identity is not None:
         raise CLIError("Cannot assign identity without enabling managed identity.")
 
+    if enable_vnet_integration:
+        if not enable_managed_identity and assign_identity is None:
+            raise CLIError("When vnet integration is enabled, either system-assigned or "
+                           "user-assigned identity must be provided.")
+
     fleet = fleet_model(
         location=location,
         tags=tags,
@@ -113,11 +119,17 @@ def create_fleet(cmd,
     )
 
     if enable_private_cluster:
+        # provider registration state being is checked to ensure that the Fleet service principal is available
+        # to create the role assignment on the subnet
+        if not is_rp_registered(cmd):
+            raise CLIError("The Microsoft.ContainerService resource provider is not registered."
+                           "Run `az provider register -n Microsoft.ContainerService --wait`.")
         assign_network_contributor_role_to_subnet(cmd, FLEET_1P_APP_ID, agent_subnet_id)
 
-    if enable_vnet_integration:
-        assign_network_contributor_role_to_subnet(cmd, get_msi_object_id(cmd, assign_identity), apiserver_subnet_id)
-        assign_network_contributor_role_to_subnet(cmd, get_msi_object_id(cmd, assign_identity), agent_subnet_id)
+    if enable_vnet_integration and assign_identity is not None:
+        object_id = get_msi_object_id(cmd, assign_identity)
+        assign_network_contributor_role_to_subnet(cmd, object_id, apiserver_subnet_id)
+        assign_network_contributor_role_to_subnet(cmd, object_id, agent_subnet_id)
 
     return sdk_no_wait(no_wait,
                        client.begin_create_or_update,