[Core] aaz: Implement wait command

src/azure-cli-core/azure/cli/core/aaz/__init__.py

@@ -15,6 +15,7 @@
 from ._arg_fmt import AAZStrArgFormat, AAZIntArgFormat, AAZFloatArgFormat, AAZBoolArgFormat, AAZObjectArgFormat, \
     AAZDictArgFormat, AAZListArgFormat, AAZResourceLocationArgFormat, AAZResourceIdArgFormat, AAZSubscriptionIdArgFormat
 from ._base import AAZValuePatch, AAZUndefined
-from ._command import AAZCommand, AAZCommandGroup, register_command, register_command_group, load_aaz_command_table
+from ._command import AAZCommand, AAZWaitCommand, AAZCommandGroup, \
+    register_command, register_command_group, load_aaz_command_table
 from ._field_type import AAZIntType, AAZFloatType, AAZStrType, AAZBoolType, AAZDictType, AAZListType, AAZObjectType
 from ._operation import AAZHttpOperation, AAZJsonInstanceUpdateOperation, AAZGenericInstanceUpdateOperation

src/azure-cli-core/azure/cli/core/aaz/_command.py

@@ -215,6 +215,26 @@ def executor_wrapper(next_link):
         return AAZPaged(executor=executor_wrapper, extract_result=extract_result)
 
 
+class AAZWaitCommand(AAZCommand):
+    """Support wait command"""
+
+    def __init__(self, loader):
+        from azure.cli.core.commands.command_operation import WaitCommandOperation
+        super().__init__(loader)
+
+        # add wait args in commands
+        for param_name, argtype in WaitCommandOperation.wait_args().items():
+            self.arguments[param_name] = argtype
+
+    def __call__(self, *args, **kwargs):
+        from azure.cli.core.commands.command_operation import WaitCommandOperation
+        return WaitCommandOperation.wait(
+            *args, **kwargs,
+            cli_ctx=self.cli_ctx,
+            getter=lambda **command_args: self._handler(command_args)
+        )
+
+
 def register_command_group(
         name, is_preview=False, is_experimental=False, hide=False, redirect=None, expiration=None):
     """This decorator is used to register an AAZCommandGroup as a cli command group.

src/azure-cli-core/azure/cli/core/commands/command_operation.py

@@ -383,13 +383,7 @@ def __init__(self, command_loader, op_path, **merged_kwargs):
 
     def handler(self, command_args):    # pylint: disable=too-many-statements, too-many-locals
         """ Callback function of CLICommand handler """
-        from msrest.exceptions import ClientException
-        from azure.core.exceptions import HttpResponseError
-        from knack.util import CLIError
-        from azure.cli.core.commands.arm import EXCLUDED_NON_CLIENT_PARAMS, verify_property
-        from azure.cli.core.commands.progress import IndeterminateProgressBar
-
-        import time
+        from azure.cli.core.commands.arm import EXCLUDED_NON_CLIENT_PARAMS
 
         op = self.get_op_handler(self.op_path)
         getter_args = dict(extract_args_from_signature(op, excluded_params=EXCLUDED_NON_CLIENT_PARAMS))
@@ -405,6 +399,18 @@ def handler(self, command_args):    # pylint: disable=too-many-statements, too-m
 
         getter = self.get_op_handler(self.op_path)      # Fetch op handler again after cmd property is set
 
+        return self.wait(command_args, cli_ctx=self.cli_ctx, getter=getter)
+
+    @classmethod
+    def wait(cls, command_args, cli_ctx, getter):
+        from msrest.exceptions import ClientException
+        from azure.core.exceptions import HttpResponseError
+        from knack.util import CLIError
+        from azure.cli.core.azclierror import InvalidArgumentValueError, AzureResponseError
+        from azure.cli.core.commands.arm import verify_property
+        from azure.cli.core.commands.progress import IndeterminateProgressBar
+        import time
+
         timeout = command_args.pop('timeout')
         interval = command_args.pop('interval')
         wait_for_created = command_args.pop('created')
@@ -414,10 +420,10 @@ def handler(self, command_args):    # pylint: disable=too-many-statements, too-m
         custom_condition = command_args.pop('custom')
         if not any([wait_for_created, wait_for_updated, wait_for_deleted,
                     wait_for_exists, custom_condition]):
-            raise CLIError(
+            raise InvalidArgumentValueError(
                 "incorrect usage: --created | --updated | --deleted | --exists | --custom JMESPATH")
 
-        progress_indicator = IndeterminateProgressBar(self.cli_ctx, message='Waiting')
+        progress_indicator = IndeterminateProgressBar(cli_ctx, message='Waiting')
         progress_indicator.begin()
         for _ in range(0, timeout, interval):
             try:
@@ -426,13 +432,13 @@ def handler(self, command_args):    # pylint: disable=too-many-statements, too-m
                 if wait_for_exists:
                     progress_indicator.end()
                     return None
-                provisioning_state = self._get_provisioning_state(instance)
+                provisioning_state = cls._get_provisioning_state(instance)
                 # until we have any needs to wait for 'Failed', let us bail out on this
                 if provisioning_state:
                     provisioning_state = provisioning_state.lower()
                 if provisioning_state == 'failed':
                     progress_indicator.stop()
-                    raise CLIError('The operation failed')
+                    raise AzureResponseError('The operation failed')
                 if ((wait_for_created or wait_for_updated) and provisioning_state == 'succeeded') or \
                         custom_condition and bool(verify_property(instance, custom_condition)):
                     progress_indicator.end()
@@ -457,26 +463,32 @@ def handler(self, command_args):    # pylint: disable=too-many-statements, too-m
 
     @staticmethod
     def _get_provisioning_state(instance):
-        provisioning_state = getattr(instance, 'provisioning_state', None)
+        from knack.util import todict
+        result = todict(instance)
+        provisioning_state = result.get('provisioning_state', result.get('provisioningState', None))
         if not provisioning_state:
             # some SDK, like resource-group, has 'provisioning_state' under 'properties'
-            properties = getattr(instance, 'properties', None)
+            properties = result.get('properties', None)
             if properties:
-                provisioning_state = getattr(properties, 'provisioning_state', None)
+                provisioning_state = properties.get('provisioning_state', properties.get('provisioningState', None))
                 # some SDK, like keyvault, has 'provisioningState' under 'properties.additional_properties'
                 if not provisioning_state:
-                    additional_properties = getattr(properties, 'additional_properties', {})
-                    provisioning_state = additional_properties.get('provisioningState')
-
-                # some SDK, like resource, has 'provisioningState' under 'properties' dict
-                if not provisioning_state:
-                    provisioning_state = properties.get('provisioningState')
+                    additional_properties = properties.get('additional_properties',
+                                                           properties.get('additionalProperties', {}))
+                    provisioning_state = additional_properties.get('provisioning_state',
+                                                                   additional_properties.get('provisioningState', None))
         return provisioning_state
 
     def arguments_loader(self):
         """ Callback function of CLICommand arguments_loader """
         cmd_args = self.load_getter_op_arguments(self.op_path)
 
+        cmd_args.update(self.wait_args())
+        return list(cmd_args.items())
+
+    @staticmethod
+    def wait_args():
+        cmd_args = {}
         group_name = 'Wait Condition'
         cmd_args['timeout'] = CLICommandArgument(
             'timeout', options_list=['--timeout'], default=3600, arg_group=group_name, type=int,
@@ -508,7 +520,7 @@ def arguments_loader(self):
                  "provisioningState!='InProgress', "
                  "instanceView.statuses[?code=='PowerState/running']"
         )
-        return list(cmd_args.items())
+        return cmd_args
 
     def description_loader(self):
         """ Callback function of CLICommand description_loader """