diff --git a/src/azure-cli/azure/cli/command_modules/storage/_params.py b/src/azure-cli/azure/cli/command_modules/storage/_params.py index 6176b6b9e22..80b10a379d9 100644 --- a/src/azure-cli/azure/cli/command_modules/storage/_params.py +++ b/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', diff --git a/src/azure-cli/azure/cli/command_modules/storage/completers.py b/src/azure-cli/azure/cli/command_modules/storage/completers.py index acc8b2427cb..56b5c5e8f4a 100644 --- a/src/azure-cli/azure/cli/command_modules/storage/completers.py +++ b/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) diff --git a/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_completers.py b/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_completers.py new file mode 100644 index 00000000000..7e71ad43a24 --- /dev/null +++ b/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() diff --git a/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_validators.py b/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_validators.py index f371b38e115..44fe3ec828a 100644 --- a/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_validators.py +++ b/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()