diff --git a/src/apic-extension/azext_apic_extension/custom.py b/src/apic-extension/azext_apic_extension/custom.py index 77d47e0dec2..2f7650d79d0 100644 --- a/src/apic-extension/azext_apic_extension/custom.py +++ b/src/apic-extension/azext_apic_extension/custom.py @@ -16,6 +16,7 @@ import yaml import requests from knack.log import get_logger +from knack.util import CLIError import chardet from azure.cli.core.aaz._arg import AAZStrArg from .command_patches import ImportAPIDefinitionExtension @@ -32,11 +33,11 @@ def pre_operations(self): # Check the size of 'value' if format is inline and raise error if value is greater than 3 mb if args.format == 'inline': - value_size_bytes = sys.getsizeof(args.value) + value_size_bytes = sys.getsizeof(str(args.value)) value_size_mb = value_size_bytes / (1024 * 1024) # Convert bytes to megabytes if value_size_mb > 3: - logger.error('The size of "value" is greater than 3 MB. ' - 'Please use --format "url" to import the specification from a URL for size greater than 3 mb.') + raise CLIError('The size of "value" is greater than 3 MB. ' + 'Please use --format "link" to import the specification from a URL for size greater than 3 mb.') class ExportSpecificationExtension(ExportAPIDefinitionExtension): diff --git a/src/apic-extension/azext_apic_extension/tests/latest/test_definition_commands.py b/src/apic-extension/azext_apic_extension/tests/latest/test_definition_commands.py index d1e2b8f7da8..f09e9c3f099 100644 --- a/src/apic-extension/azext_apic_extension/tests/latest/test_definition_commands.py +++ b/src/apic-extension/azext_apic_extension/tests/latest/test_definition_commands.py @@ -7,6 +7,7 @@ import json import os +from knack.util import CLIError from azure.cli.testsdk import ScenarioTest, ResourceGroupPreparer from .utils import ApicServicePreparer, ApicApiPreparer, ApicVersionPreparer, ApicDefinitionPreparer @@ -185,4 +186,24 @@ def test_definition_import_from_file(self): assert exported_content == imported_content, "The exported content is not the same as the imported content." finally: - os.remove(exported_file_path) \ No newline at end of file + os.remove(exported_file_path) + + @ResourceGroupPreparer(name_prefix="clirg", location='eastus', random_name_length=32) + @ApicServicePreparer() + @ApicApiPreparer() + @ApicVersionPreparer() + @ApicDefinitionPreparer() + def test_definition_import_large_value(self): + self.kwargs.update({ + 'specification': '{"name":"openapi","version":"3.0.0"}', + 'file_name': "test_definition_import_large_value.txt" + }) + + try: + with open(self.kwargs['file_name'], 'w') as file: + file.write('a' * 4 * 1024 * 1024) # generate a 4MB file + + with self.assertRaisesRegexp(CLIError, 'The size of "value" is greater than 3 MB. Please use --format "link" to import the specification from a URL for size greater than 3 mb.') as cm: + self.cmd('az apic api definition import-specification -g {rg} -s {s} --api-id {api} --version-id {v} --definition-id {d} --format "inline" --specification \'{specification}\' --value "@{file_name}"') + finally: + os.remove(self.kwargs['file_name'])