[Storage] Reduce debug-time import failures in az storage container list

src/azure-cli/azure/cli/command_modules/storage/_params.py

@@ -3,6 +3,8 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+from importlib import import_module
+
 from azure.cli.core.profiles import ResourceType
 from azure.cli.core.commands.validators import get_default_location_from_resource_group
 from azure.cli.core.commands.parameters import (tags_type, file_type, get_location_type,
@@ -27,6 +29,13 @@
                           SshPublicKeyAddAction, user_delegation_oid_validator, user_delegation_tid_validator)
 
 
+def _get_legacy_storage_data_plane_sdk(module_name, class_name):
+    try:
+        return getattr(import_module(module_name), class_name)
+    except (ImportError, AttributeError):
+        return None
+
+
 def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statements, too-many-lines, too-many-branches, line-too-long
     from argcomplete.completers import FilesCompleter
 
@@ -36,8 +45,8 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
 
     from .completers import get_storage_name_completion_list
 
-    t_base_blob_service = self.get_sdk('blob.baseblobservice#BaseBlobService')
-    t_file_service = self.get_sdk('file#FileService')
+    t_base_blob_service = _get_legacy_storage_data_plane_sdk('azure.storage.blob.baseblobservice', 'BaseBlobService')
+    t_file_service = _get_legacy_storage_data_plane_sdk('azure.storage.file', 'FileService')
     t_share_service = self.get_sdk('_share_service_client#ShareServiceClient',
                                    resource_type=ResourceType.DATA_STORAGE_FILESHARE)
     t_queue_service = self.get_sdk('_queue_service_client#QueueServiceClient',
@@ -2213,7 +2222,7 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
         c.register_path_argument()
         c.register_sas_arguments()
         c.extra('share_name', share_name_type, required=True)
-        t_file_svc = self.get_sdk('file.fileservice#FileService')
+        t_file_svc = _get_legacy_storage_data_plane_sdk('azure.storage.file.fileservice', 'FileService')
         t_file_permissions = self.get_sdk('_models#FileSasPermissions',
                                           resource_type=ResourceType.DATA_STORAGE_FILESHARE)
         c.argument('id', options_list='--policy-name',

src/azure-cli/azure/cli/command_modules/storage/completers.py

@@ -9,6 +9,14 @@
 from ._validators import validate_client_parameters
 
 
+def _get_legacy_file_service_class():
+    try:
+        from azure.storage.file import FileService
+        return FileService
+    except ImportError:
+        return None
+
+
 @Completer
 def file_path_completer(cmd, prefix, namespace):
     from azure.common import AzureMissingResourceHttpError
@@ -18,7 +26,9 @@ def file_path_completer(cmd, prefix, namespace):
 
     validate_client_parameters(cmd, namespace)
 
-    t_file_service = cmd.get_models('file#FileService')
+    t_file_service = _get_legacy_file_service_class()
+    if t_file_service is None:
+        return []
     client = get_storage_client(cmd.cli_ctx, t_file_service, namespace)
 
     share_name = namespace.share_name
@@ -50,7 +60,9 @@ def dir_path_completer(cmd, prefix, namespace):
 
     validate_client_parameters(cmd, namespace)
 
-    t_file_service = cmd.get_models('file#FileService')
+    t_file_service = _get_legacy_file_service_class()
+    if t_file_service is None:
+        return []
     client = get_storage_client(cmd.cli_ctx, t_file_service, namespace)
 
     share_name = namespace.share_name
@@ -75,6 +87,8 @@ def dir_path_completer(cmd, prefix, namespace):
 def get_storage_name_completion_list(service, func, parent=None):
     @Completer
     def completer(cmd, _, namespace):
+        if service is None:
+            return []
         validate_client_parameters(cmd, namespace)
         client = get_storage_client(cmd.cli_ctx, service, namespace)
         if parent:
@@ -91,6 +105,8 @@ def completer(cmd, _, namespace):
 def get_storage_acl_name_completion_list(service, container_param, func):
     @Completer
     def completer(cmd, _, namespace):
+        if service is None:
+            return []
         validate_client_parameters(cmd, namespace)
         client = get_storage_client(cmd.cli_ctx, service, namespace)
         container_name = getattr(namespace, container_param)

src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_completers.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import unittest
+from argparse import Namespace
+from unittest import mock
+
+from azure.cli.command_modules.storage.completers import (get_storage_name_completion_list,
+                                                          get_storage_acl_name_completion_list,
+                                                          file_path_completer)
+from azure.cli.command_modules.storage.tests.latest.test_storage_validators import MockCLI, MockCmd
+
+
+class TestStorageCompleters(unittest.TestCase):
+
+    def setUp(self):
+        self.cli = MockCLI()
+
+    def test_storage_name_completer_returns_empty_with_missing_legacy_service(self):
+        completer = get_storage_name_completion_list(None, 'list_containers')
+        self.assertEqual(completer.func(MockCmd(self.cli), '', Namespace()), [])
+
+    def test_storage_acl_name_completer_returns_empty_with_missing_legacy_service(self):
+        completer = get_storage_acl_name_completion_list(None, 'container_name', 'get_container_acl')
+        self.assertEqual(completer.func(MockCmd(self.cli), '', Namespace(container_name='container')), [])
+
+    @mock.patch('azure.cli.command_modules.storage.completers.validate_client_parameters')
+    @mock.patch('azure.cli.command_modules.storage.completers._get_legacy_file_service_class', return_value=None)
+    def test_file_path_completer_returns_empty_with_missing_legacy_service(self, _, __):
+        self.assertEqual(file_path_completer.func(MockCmd(self.cli), '', Namespace(share_name='share')), [])
+
+
+if __name__ == '__main__':
+    unittest.main()

src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_validators.py

@@ -21,8 +21,6 @@
                                                            validate_encryption_source, validate_source_uri,
                                                            validate_encryption_services, as_user_validator,
                                                            get_not_none_validator)
-
-
 class MockCLI(CLI):
     def __init__(self):
         super().__init__(cli_name='mock_cli', config_dir=GLOBAL_CONFIG_DIR,
@@ -183,7 +181,6 @@ def test_get_not_none_validator(self):
         validate_arg(cmd, Namespace(arg=0))
         validate_arg(cmd, Namespace(arg=False))
 
-
 class TestEncryptionValidators(unittest.TestCase):
     def setUp(self):
         self.cli = MockCLI()