refine client

Juliehzl · Juliehzl · commit 10d11c049073 · 2021-04-21T14:22:12.000+08:00
diff --git a/src/azure-cli/azure/cli/command_modules/batchai/custom.py b/src/azure-cli/azure/cli/command_modules/batchai/custom.py
@@ -143,8 +143,8 @@ def _ensure_subnet_is_valid(client, subnet, nfs_resource_group, nfs_workspace, n
 
 
 def _get_storage_management_client(cli_ctx):
-    from azure.mgmt.storage import StorageManagementClient
-    return get_mgmt_service_client(cli_ctx, StorageManagementClient)
+    from azure.cli.core.profiles._shared import ResourceType
+    return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_STORAGE)
 
 
 def _get_storage_account_key(cli_ctx, account_name, account_key):
