[Core] aaz: Fix early exit for long running commands run with --no-wait parameter

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

@@ -127,7 +127,12 @@ def __call__(self, *args, **kwargs):
 
     def _handler(self, command_args):
         # command_args will be parsed by AAZCommandCtx
-        self.ctx = AAZCommandCtx(cli_ctx=self.cli_ctx, schema=self.get_arguments_schema(), command_args=command_args)
+        self.ctx = AAZCommandCtx(
+            cli_ctx=self.cli_ctx,
+            schema=self.get_arguments_schema(),
+            command_args=command_args,
+            no_wait_arg='no_wait' if self.supports_no_wait else None,
+        )
         self.ctx.format_args()
 
     def _cli_arguments_loader(self):
@@ -190,11 +195,15 @@ def processor(schema, result):
 
         return value.to_serialized_data(processor=processor)
 
-    @staticmethod
-    def build_lro_poller(executor, extract_result):
+    def build_lro_poller(self, executor, extract_result):
         """ Build AAZLROPoller instance to support long running operation
         """
-        return AAZLROPoller(polling_generator=executor, result_callback=extract_result)
+        polling_generator = executor()
+        if self.ctx.lro_no_wait:
+            # run until yield the first polling
+            _ = next(polling_generator)
+            return None
+        return AAZLROPoller(polling_generator=polling_generator, result_callback=extract_result)
 
     def build_paging(self, executor, extract_result):
         """ Build AAZPaged instance to support paging

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

@@ -15,7 +15,7 @@
 
 class AAZCommandCtx:
 
-    def __init__(self, cli_ctx, schema, command_args):
+    def __init__(self, cli_ctx, schema, command_args, no_wait_arg=None):
         self._cli_ctx = cli_ctx
         self._profile = Profile(cli_ctx=cli_ctx)
         self._subscription_id = None
@@ -31,8 +31,10 @@ def __init__(self, cli_ctx, schema, command_args):
         self._vars_schema = AAZObjectType()
         self.vars = AAZObject(schema=self._vars_schema, data={})
         self.generic_update_args = command_args.get(AAZGenericUpdateAction.DEST, None)
-
-        self.next_link = AAZUndefined  # support paging
+        # support no wait
+        self.lro_no_wait = command_args.get(no_wait_arg, False) if no_wait_arg else False
+        # support paging
+        self.next_link = AAZUndefined
 
     def format_args(self):
         # TODO: apply format for argument values

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

@@ -77,7 +77,7 @@ def _start(self):
          the operation.
         """
         try:
-            for polling_method in self._polling_generator():
+            for polling_method in self._polling_generator:
                 self._polling_method = polling_method
                 try:
                     self._polling_method.run()