refine client

Azure · Apr 21, 2021 · 10d11c0 · 10d11c0
1 parent 8bceba4
commit 10d11c0
Showing 1 changed file with 2 additions and 2 deletions.
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):