Digital Twin wait commands

azext_iot/digitaltwins/_help.py

@@ -113,6 +113,22 @@ def load_digitaltwins_help():
             az dt delete -n {instance_name} -y --no-wait
     """
 
+    helps["dt wait"] = """
+        type: command
+        short-summary: Wait until an operation on an Digital Twins instance is complete.
+
+        examples:
+        - name: Wait until an arbitrary instance is created.
+          text: >
+            az dt wait -n {instance_name} --created
+        - name: Wait until an existing instance is deleted.
+          text: >
+            az dt wait -n {instance_name} --deleted
+        - name: Wait until an existing instance's publicNetworkAccess property is set to Enabled
+          text: >
+            az dt wait -n {instance_name} --custom "publicNetworkAccess=='Enabled'"
+    """
+
     helps["dt reset"] = """
         type: command
         short-summary: Reset an existing Digital Twins instance by deleting associated
@@ -232,6 +248,22 @@ def load_digitaltwins_help():
             az dt endpoint delete -n {instance_name} --endpoint-name {endpoint_name} -y --no-wait
     """
 
+    helps["dt endpoint wait"] = """
+        type: command
+        short-summary: Wait until an endpoint operation is done.
+
+        examples:
+        - name: Wait until an endpoint for an instance is created.
+          text: >
+            az dt endpoint wait -n {instance_name} --endpoint-name {endpoint_name} --created
+        - name: Wait until an existing endpoint is deleted from an instance.
+          text: >
+            az dt endpoint wait -n {instance_name} --endpoint-name {endpoint_name} --deleted
+        - name: Wait until an existing endpoint's primaryConnectionString is null.
+          text: >
+            az dt endpoint wait -n {instance_name} --endpoint-name {endpoint_name} --custom "properties.primaryConnectionString==null"
+    """
+
     helps["dt network"] = """
         type: group
         short-summary: Manage Digital Twins network configuration including private links and endpoint connections.
@@ -301,7 +333,6 @@ def load_digitaltwins_help():
         - name: Approve a pending private-endpoint connection associated with the instance and add a description.
           text: >
             az dt network private-endpoint connection set -n {instance_name} --cn {connection_name} --status Approved --desc "A description."
-
         - name: Reject a private-endpoint connection associated with the instance and add a description.
           text: >
             az dt network private-endpoint connection set -n {instance_name} --cn {connection_name} --status Rejected --desc "Does not comply."
@@ -321,6 +352,23 @@ def load_digitaltwins_help():
             az dt network private-endpoint connection delete -n {instance_name} --cn ba8408b6-1372-41b2-aef8-af43afc4729f -y --no-wait
     """
 
+    helps["dt network private-endpoint connection wait"] = """
+        type: command
+        short-summary: Wait until an operation on a private-endpoint connection is complete.
+
+        examples:
+        - name: Wait until the existing private-endpoint connection named ba8408b6-1372-41b2-aef8-af43afc4729f state is updated.
+          text: >
+            az dt network private-endpoint connection wait -n {instance_name} --cn ba8408b6-1372-41b2-aef8-af43afc4729f --updated
+
+        - name: Wait until the existing private-endpoint connection named ba8408b6-1372-41b2-aef8-af43afc4729f is deleted.
+          text: >
+            az dt network private-endpoint connection wait -n {instance_name} --cn ba8408b6-1372-41b2-aef8-af43afc4729f --deleted
+        - name: Wait until the existing private-endpoint connection named ba8408b6-1372-41b2-aef8-af43afc4729f has no actions required in the privateLinkServiceConnectionState property.
+          text: >
+            az dt network private-endpoint connection wait -n {instance_name} --cn ba8408b6-1372-41b2-aef8-af43afc4729f --custom "properties.privateLinkServiceConnectionState.actionsRequired=='None'"
+    """
+
     helps["dt role-assignment"] = """
         type: group
         short-summary: Manage RBAC role assignments for a Digital Twins instance.

azext_iot/digitaltwins/command_map.py

@@ -42,6 +42,7 @@ def load_digitaltwins_commands(self, _):
         cmd_group.show_command("show", "show_instance")
         cmd_group.command("list", "list_instances")
         cmd_group.command("delete", "delete_instance", confirmation=True, supports_no_wait=True)
+        cmd_group.wait_command("wait", "wait_instance")
         cmd_group.command("reset", "reset_instance", confirmation=True, is_preview=True)
 
     with self.command_group(
@@ -64,6 +65,9 @@ def load_digitaltwins_commands(self, _):
             ),
         )
         cmd_group.command("delete", "delete_endpoint", confirmation=True, supports_no_wait=True)
+        cmd_group.wait_command(
+            "wait", "wait_endpoint"
+        )
 
     with self.command_group(
         "dt endpoint create", command_type=digitaltwins_resource_ops
@@ -169,3 +173,6 @@ def load_digitaltwins_commands(self, _):
         cmd_group.show_command("show", "show_private_endpoint_conn")
         cmd_group.command("list", "list_private_endpoint_conns")
         cmd_group.command("delete", "delete_private_endpoint_conn", confirmation=True, supports_no_wait=True)
+        cmd_group.wait_command(
+            "wait", "wait_private_endpoint_conn"
+        )

azext_iot/digitaltwins/commands_resource.py

@@ -59,6 +59,11 @@ def delete_instance(cmd, name, resource_group_name=None):
     return rp.delete(name=name, resource_group_name=resource_group_name)
 
 
+def wait_instance(cmd, name, resource_group_name=None):
+    rp = ResourceProvider(cmd)
+    return rp.find_instance(name=name, resource_group_name=resource_group_name, wait=True)
+
+
 def reset_instance(cmd, name, resource_group_name=None):
     delete_all_models(cmd, name, resource_group_name)
     delete_all_twin(cmd, name, resource_group_name)
@@ -83,6 +88,16 @@ def delete_endpoint(cmd, name, endpoint_name, resource_group_name=None):
     )
 
 
+def wait_endpoint(cmd, name, endpoint_name, resource_group_name=None):
+    rp = ResourceProvider(cmd)
+    return rp.get_endpoint(
+        name=name,
+        endpoint_name=endpoint_name,
+        resource_group_name=resource_group_name,
+        wait=True
+    )
+
+
 def add_endpoint_eventgrid(
     cmd,
     name,
@@ -225,3 +240,13 @@ def delete_private_endpoint_conn(cmd, name, conn_name, resource_group_name=None)
     return rp.delete_private_endpoint_conn(
         name=name, resource_group_name=resource_group_name, conn_name=conn_name
     )
+
+
+def wait_private_endpoint_conn(cmd, name, conn_name, resource_group_name=None):
+    rp = ResourceProvider(cmd)
+    return rp.get_private_endpoint_conn(
+        name=name,
+        resource_group_name=resource_group_name,
+        conn_name=conn_name,
+        wait=True
+    )

azext_iot/digitaltwins/params.py

@@ -142,6 +142,9 @@ def load_digitaltwins_arguments(self, _):
             help="Role name or Id the system assigned identity will have.",
         )
 
+    with self.argument_context("dt wait") as context:
+        context.ignore("updated")
+
     with self.argument_context("dt endpoint create") as context:
         context.argument(
             "dead_letter_secret",
@@ -249,6 +252,9 @@ def load_digitaltwins_arguments(self, _):
             arg_group="Service Bus Topic",
         )
 
+    with self.argument_context("dt endpoint wait") as context:
+        context.ignore("updated")
+
     with self.argument_context("dt twin") as context:
         context.argument(
             "query_command",
@@ -423,3 +429,7 @@ def load_digitaltwins_arguments(self, _):
             help="A message indicating if changes on the service provider require any updates on the consumer.",
             arg_group="Private-Endpoint",
         )
+
+    with self.argument_context("dt network private-endpoint connection wait") as context:
+        context.ignore("created")
+        context.ignore("exists")

azext_iot/digitaltwins/providers/resource.py

@@ -117,20 +117,22 @@ def list_by_resouce_group(self, resource_group_name):
         except ErrorResponseException as e:
             raise CLIError(unpack_msrest_error(e))
 
-    def get(self, name, resource_group_name):
+    def get(self, name, resource_group_name, wait=False):
         try:
             return self.mgmt_sdk.digital_twins.get(
                 resource_name=name, resource_group_name=resource_group_name
             )
         except ErrorResponseException as e:
+            if wait:
+                e.status_code = e.response.status_code
+                raise e
             raise CLIError(unpack_msrest_error(e))
 
-    def find_instance(self, name, resource_group_name=None):
+    def find_instance(self, name, resource_group_name=None, wait=False):
         if resource_group_name:
-            try:
-                return self.get(name=name, resource_group_name=resource_group_name)
-            except ErrorResponseException as e:
-                raise CLIError(unpack_msrest_error(e))
+            return self.get(
+                name=name, resource_group_name=resource_group_name, wait=wait
+            )
 
         dt_collection_pager = self.list()
         dt_collection = []
@@ -221,7 +223,7 @@ def remove_role(self, name, assignee, role_type=None, resource_group_name=None):
 
     # Endpoints
 
-    def get_endpoint(self, name, endpoint_name, resource_group_name=None):
+    def get_endpoint(self, name, endpoint_name, resource_group_name=None, wait=False):
         target_instance = self.find_instance(
             name=name, resource_group_name=resource_group_name
         )
@@ -235,6 +237,9 @@ def get_endpoint(self, name, endpoint_name, resource_group_name=None):
                 resource_group_name=resource_group_name,
             )
         except ErrorResponseException as e:
+            if wait:
+                e.status_code = e.response.status_code
+                raise e
             raise CLIError(unpack_msrest_error(e))
 
     def list_endpoints(self, name, resource_group_name=None):
@@ -417,7 +422,13 @@ def set_private_endpoint_conn(
         except ErrorResponseException as e:
             raise CLIError(unpack_msrest_error(e))
 
-    def get_private_endpoint_conn(self, name, conn_name, resource_group_name=None):
+    def get_private_endpoint_conn(
+        self,
+        name,
+        conn_name,
+        resource_group_name=None,
+        wait=False
+    ):
         target_instance = self.find_instance(
             name=name, resource_group_name=resource_group_name
         )
@@ -428,10 +439,12 @@ def get_private_endpoint_conn(self, name, conn_name, resource_group_name=None):
             return self.mgmt_sdk.private_endpoint_connections.get(
                 resource_group_name=resource_group_name,
                 resource_name=name,
-                private_endpoint_connection_name=conn_name,
-                raw=True,
-            ).response.json()
+                private_endpoint_connection_name=conn_name
+            )
         except ErrorResponseException as e:
+            if wait:
+                e.status_code = e.response.status_code
+                raise e
             raise CLIError(unpack_msrest_error(e))
 
     def list_private_endpoint_conns(self, name, resource_group_name=None):

azext_iot/tests/digitaltwins/__init__.py

@@ -160,23 +160,22 @@ def tearDown(self):
 
     # Needed because the DT service will indicate provisioning is finished before it actually is.
     def wait_for_hostname(
-        self, instance: dict, wait_in_sec: int = 10, interval: int = 4
+        self, instance: dict, wait_in_sec: int = 10, interval: int = 7
     ):
         from time import sleep
-
-        while interval >= 1:
-            logger.info(
-                "Waiting :{} (sec) for provisioning to complete.".format(wait_in_sec)
+        sleep(wait_in_sec)
+
+        self.embedded_cli.invoke(
+            "dt wait -n {} -g {} --custom \"hostName && provisioningState=='Succeeded'\" --interval {} --timeout {}".format(
+                instance["name"],
+                instance["resourceGroup"],
+                wait_in_sec,
+                wait_in_sec * interval
             )
-            sleep(wait_in_sec)
-            interval = interval - 1
-            refereshed_instance = self.embedded_cli.invoke(
-                "dt show -n {} -g {}".format(
-                    instance["name"], instance["resourceGroup"]
-                )
-            ).as_json()
-
-            if refereshed_instance.get("hostName") and refereshed_instance["provisioningState"] == "Succeeded":
-                return refereshed_instance
-
-        return instance
+        )
+        refereshed_instance = self.embedded_cli.invoke(
+            "dt show -n {} -g {}".format(
+                instance["name"], instance["resourceGroup"]
+            )
+        ).as_json()
+        return refereshed_instance if refereshed_instance else instance

azext_iot/tests/digitaltwins/test_dt_privatelinks_lifecycle_int.py

@@ -158,6 +158,26 @@ def test_dt_privatelinks(self):
                 instance_name, self.rg, instance_connection_id, random_desc_rejected
             )
         ).get_output_in_json()
+
+        assert (
+            set_connection_output["properties"]["privateLinkServiceConnectionState"]["status"]
+            == "Rejected"
+        )
+        assert (
+            set_connection_output["properties"]["privateLinkServiceConnectionState"]["description"]
+            == random_desc_rejected
+        )
+
+        self.cmd("dt network private-endpoint connection wait -n {} -g {} --cn {} --updated --interval 1 --timeout 30".format(
+            instance_name, self.rg, instance_connection_id
+        ))
+
+        set_connection_output = self.cmd(
+            "dt network private-endpoint connection show -n {} -g {} --cn {}".format(
+                instance_name, self.rg, instance_connection_id
+            )
+        ).get_output_in_json()
+
         assert (
             set_connection_output["properties"]["privateLinkServiceConnectionState"]["status"]
             == "Rejected"
@@ -173,6 +193,12 @@ def test_dt_privatelinks(self):
             )
         )
 
+        self.cmd(
+            "dt network private-endpoint connection wait -n {} -g {} --cn {} --deleted --interval 1 --timeout 30".format(
+                instance_name, self.rg, instance_connection_id
+            )
+        )
+
         list_priv_endpoints = self.cmd(
             "dt network private-endpoint connection list -n {} -g {}".format(
                 instance_name,