refactor: re-add cluster lookup to delete, show tree visual after init.

azext_edge/edge/_help.py

@@ -560,8 +560,12 @@ def load_iotops_help():
     ] = """
         type: command
         short-summary: Delete IoT Operations from the cluster.
-        long-summary: The operation uses Azure Resource Graph to determine correlated resources.
-          Resource Graph being eventually consistent does not guarantee a synchronized state at the time of execution.
+        long-summary: |
+            The name of either the instance or cluster must be provided.
+
+            The operation uses Azure Resource Graph to determine correlated resources.
+            Resource Graph being eventually consistent does not guarantee a synchronized state at the
+            time of execution.
 
         examples:
         - name: Minimum input for complete deletion.
@@ -573,6 +577,9 @@ def load_iotops_help():
         - name: Force deletion regardless of warnings. May lead to errors.
           text: >
             az iot ops delete -n myinstance -g myresourcegroup --force
+        - name: Use cluster name instead of instance for lookup.
+          text: >
+            az iot ops delete --cluster mycluster -g myresourcegroup
         - name: Reverse application of init.
           text: >
             az iot ops delete -n myinstance -g myresourcegroup --include-deps

azext_edge/edge/commands_edge.py

@@ -216,7 +216,8 @@ def create_instance(
 def delete(
     cmd,
     resource_group_name: str,
-    instance_name: str,
+    instance_name: Optional[str] = None,
+    cluster_name: Optional[str] = None,
     confirm_yes: Optional[bool] = None,
     no_progress: Optional[bool] = None,
     force: Optional[bool] = None,
@@ -227,6 +228,7 @@ def delete(
     return delete_ops_resources(
         cmd=cmd,
         instance_name=instance_name,
+        cluster_name=cluster_name,
         resource_group_name=resource_group_name,
         confirm_yes=confirm_yes,
         no_progress=no_progress,

azext_edge/edge/params.py

@@ -533,6 +533,11 @@ def load_iotops_arguments(self, _):
             help="Indicates the command should remove IoT Operations dependencies. "
             "This option is intended to reverse the application of init.",
         )
+        context.argument(
+            "cluster_name",
+            options_list=["--cluster"],
+            help="Target cluster name for IoT Operations deletion.",
+        )
 
     with self.argument_context("iot ops asset") as context:
         context.argument(

azext_edge/edge/providers/orchestration/deletion.py

@@ -8,6 +8,7 @@
 from time import sleep
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
+from azure.cli.core.azclierror import ArgumentUsageError
 from knack.log import get_logger
 from rich import print
 from rich.console import NewLine
@@ -17,8 +18,8 @@
 
 from ...util.az_client import get_resource_client, wait_for_terminal_states
 from ...util.common import should_continue_prompt
+from .resource_map import IoTOperationsResource, IoTOperationsResourceMap
 from .resources import Instances
-from .resource_map import IoTOperationsResource
 
 logger = get_logger(__name__)
 
@@ -29,8 +30,9 @@
 
 def delete_ops_resources(
     cmd,
-    instance_name: str,
     resource_group_name: str,
+    instance_name: Optional[str] = None,
+    cluster_name: Optional[str] = None,
     confirm_yes: Optional[bool] = None,
     no_progress: Optional[bool] = None,
     force: Optional[bool] = None,
@@ -39,6 +41,7 @@ def delete_ops_resources(
     manager = DeletionManager(
         cmd=cmd,
         instance_name=instance_name,
+        cluster_name=cluster_name,
         resource_group_name=resource_group_name,
         no_progress=no_progress,
         include_dependencies=include_dependencies,
@@ -50,15 +53,17 @@ class DeletionManager:
     def __init__(
         self,
         cmd,
-        instance_name: str,
         resource_group_name: str,
+        instance_name: Optional[str] = None,
+        cluster_name: Optional[str] = None,
         include_dependencies: Optional[bool] = None,
         no_progress: Optional[bool] = None,
     ):
         from azure.cli.core.commands.client_factory import get_subscription_id
 
         self.cmd = cmd
         self.instance_name = instance_name
+        self.cluster_name = cluster_name
         self.resource_group_name = resource_group_name
         self.instances = Instances(self.cmd)
         self.include_dependencies = include_dependencies
@@ -77,8 +82,7 @@ def __init__(
         self._progress_shown = False
 
     def do_work(self, confirm_yes: Optional[bool] = None, force: Optional[bool] = None):
-        self.instance = self.instances.show(name=self.instance_name, resource_group_name=self.resource_group_name)
-        self.resource_map = self.instances.get_resource_map(self.instance)
+        self.resource_map = self._get_resource_map()
         # Ensure cluster exists with existing resource_map pattern.
         self.resource_map.connected_cluster.resource
         self.resource_map.refresh_resource_state()
@@ -90,9 +94,24 @@ def do_work(self, confirm_yes: Optional[bool] = None, force: Optional[bool] = No
 
         self._process(force=force)
 
+    def _get_resource_map(self) -> IoTOperationsResourceMap:
+        if not any([self.cluster_name, self.instance_name]):
+            raise ArgumentUsageError("Please provide either an instance name or cluster name.")
+
+        if self.instance_name:
+            self.instance = self.instances.show(name=self.instance_name, resource_group_name=self.resource_group_name)
+            return self.instances.get_resource_map(self.instance)
+
+        return IoTOperationsResourceMap(
+            cmd=self.cmd,
+            cluster_name=self.cluster_name,
+            resource_group_name=self.resource_group_name,
+            defer_refresh=True,
+        )
+
     def _display_resource_tree(self):
         if self._render_progress:
-            print(self.resource_map.build_tree(hide_extensions=not self.include_dependencies))
+            print(self.resource_map.build_tree(hide_extensions=not self.include_dependencies, category_color="red"))
 
     def _render_display(self, description: str):
         if self._render_progress:

azext_edge/edge/providers/orchestration/resource_map.py

@@ -127,16 +127,16 @@ def refresh_resource_state(self):
 
         self._cluster_container = refreshed_cluster_container
 
-    def build_tree(self, hide_extensions: bool = False, category_color: str = "red") -> Tree:
+    def build_tree(self, hide_extensions: bool = False, category_color: str = "cyan") -> Tree:
         tree = Tree(f"[green]{self.connected_cluster.cluster_name}")
 
         if not hide_extensions:
             extensions_node = tree.add(label=f"[{category_color}]extensions")
             [extensions_node.add(ext.display_name) for ext in self.extensions]
 
-        root_cl_node = tree.add(label=f"[{category_color}]customLocations")
         custom_locations = self.custom_locations
         if custom_locations:
+            root_cl_node = tree.add(label=f"[{category_color}]customLocations")
             for cl in custom_locations:
                 cl_node = root_cl_node.add(cl.display_name)
                 resource_sync_rules = self.get_resource_sync_rules(cl.resource_id)

azext_edge/edge/providers/orchestration/work.py

@@ -13,6 +13,7 @@
 from azure.cli.core.azclierror import AzureResponseError, ValidationError
 from azure.core.exceptions import HttpResponseError
 from knack.log import get_logger
+from rich import print
 from rich.console import NewLine
 from rich.live import Live
 from rich.padding import Padding
@@ -25,9 +26,8 @@
     get_resource_client,
     wait_for_terminal_state,
 )
-from .connected_cluster import ConnectedCluster
 from .permissions import ROLE_DEF_FORMAT_STR, PermissionManager
-from .resources import Instances
+from .resource_map import IoTOperationsResourceMap
 from .targets import InitTargets
 
 logger = get_logger(__name__)
@@ -95,7 +95,6 @@ def __init__(self, cmd):
         self.subscription_id: str = get_subscription_id(cli_ctx=cmd.cli_ctx)
         self.resource_client = get_resource_client(subscription_id=self.subscription_id)
         self.permission_manager = PermissionManager(subscription_id=self.subscription_id)
-        self.instances = Instances(cmd=cmd)
 
     def _bootstrap_ux(self, show_progress: bool = False):
         self._display = WorkDisplay()
@@ -154,13 +153,8 @@ def _build_display(self):
                 self._format_instance_desc(),
             )
 
-    def _process_connected_cluster(self) -> ConnectedCluster:
-        connected_cluster = ConnectedCluster(
-            cmd=self.cmd,
-            subscription_id=self.subscription_id,
-            cluster_name=self._targets.cluster_name,
-            resource_group_name=self._targets.resource_group_name,
-        )
+    def _process_connected_cluster(self):
+        connected_cluster = self._resource_map.connected_cluster
         cluster = connected_cluster.resource
         cluster_properties: Dict[str, Union[str, dict]] = cluster["properties"]
         cluster_validation_tuples = [
@@ -177,8 +171,6 @@ def _process_connected_cluster(self) -> ConnectedCluster:
         if self._targets.enable_fault_tolerance and cluster_properties["totalNodeCount"] < 3:
             raise ValidationError("Arc Container Storage fault tolerance enablement requires at least 3 nodes.")
 
-        return connected_cluster
-
     def _deploy_template(
         self,
         content: dict,
@@ -224,6 +216,12 @@ def execute_ops_init(
         self._active_step: int = 0
         self._targets = InitTargets(subscription_id=self.subscription_id, **kwargs)
         self._extension_map = None
+        self._resource_map = IoTOperationsResourceMap(
+            cmd=self.cmd,
+            cluster_name=self._targets.cluster_name,
+            resource_group_name=self._targets.resource_group_name,
+            defer_refresh=True,
+        )
         self._build_display()
 
         return self._do_work()
@@ -243,7 +241,7 @@ def _do_work(self):  # noqa: C901
             # Ensure connection to ARM if needed. Show remediation error message otherwise.
             self.render_display()
             verify_cli_client_connections()
-            connected_cluster = self._process_connected_cluster()
+            self._process_connected_cluster()
 
             # Pre-Flight workflow
             if self._pre_flight:
@@ -261,7 +259,7 @@ def _do_work(self):  # noqa: C901
                 if False:
                     verify_custom_locations_enabled(self.cmd)
                     verify_custom_location_namespace(
-                        connected_cluster=connected_cluster,
+                        connected_cluster=self._resource_map.connected_cluster,
                         custom_location_name=self._targets.custom_location_name,
                         namespace=self._targets.cluster_namespace,
                     )
@@ -316,7 +314,7 @@ def _do_work(self):  # noqa: C901
                 self.render_display(category=WorkCategoryKey.ENABLE_IOT_OPS)
                 _ = wait_for_terminal_state(enablement_poller)
 
-                self._extension_map = connected_cluster.get_extensions_by_type(
+                self._extension_map = self._resource_map.connected_cluster.get_extensions_by_type(
                     IOT_OPS_EXTENSION_TYPE, IOT_OPS_PLAT_EXTENSION_TYPE, SECRET_SYNC_EXTENSION_TYPE
                 )
                 self.permission_manager.apply_role_assignment(
@@ -330,16 +328,19 @@ def _do_work(self):  # noqa: C901
                     category=WorkCategoryKey.ENABLE_IOT_OPS, completed_step=WorkStepKey.DEPLOY_ENABLEMENT
                 )
 
-                # TODO @digimaun
                 if self._show_progress:
+                    self._resource_map.refresh_resource_state()
+                    resource_tree = self._resource_map.build_tree()
                     self.stop_display()
+                    print(resource_tree)
                     return
+                # TODO @digimaun - work_kpis
                 return work_kpis
 
             # Deploy IoT Ops workflow
             if self._targets.instance_name:
                 if not self._extension_map:
-                    self._extension_map = connected_cluster.get_extensions_by_type(
+                    self._extension_map = self._resource_map.connected_cluster.get_extensions_by_type(
                         IOT_OPS_EXTENSION_TYPE, IOT_OPS_PLAT_EXTENSION_TYPE, SECRET_SYNC_EXTENSION_TYPE
                     )
                     # TODO - @digmaun revisit
@@ -392,12 +393,10 @@ def _do_work(self):  # noqa: C901
                 )
 
                 if self._show_progress:
+                    self._resource_map.refresh_resource_state()
+                    resource_tree = self._resource_map.build_tree()
                     self.stop_display()
-                    self.instances.show(
-                        name=self._targets.instance_name,
-                        resource_group_name=self._targets.resource_group_name,
-                        show_tree=True,
-                    )
+                    print(resource_tree)
                     return
                 # TODO @digimaun - work_kpis
                 return work_kpis

azext_edge/tests/edge/orchestration/test_resource_map_unit.py

@@ -69,6 +69,7 @@ def _assemble_connected_cluster_mock(
 )
 @pytest.mark.parametrize("expected_resources", [None, _generate_records(5)])
 @pytest.mark.parametrize("expected_resource_sync_rules", [None, _generate_records()])
+@pytest.mark.parametrize("category_color", [None, "red"])
 def test_resource_map(
     mocker,
     mocked_cmd: Mock,
@@ -77,6 +78,7 @@ def test_resource_map(
     expected_custom_locations: Optional[List[dict]],
     expected_resources: Optional[List[dict]],
     expected_resource_sync_rules: Optional[List[dict]],
+    category_color: Optional[str],
 ):
     from azext_edge.edge.providers.orchestration.resource_map import (
         IoTOperationsResourceMap,
@@ -117,13 +119,18 @@ def test_resource_map(
             )
             _assert_ops_resource_eq(resource_map.get_resource_sync_rules(cl["id"]), expected_resource_sync_rules)
 
+    kwargs = {}
+    if category_color:
+        kwargs["category_color"] = category_color
+
     _assert_tree(
-        resource_map.build_tree(),
+        resource_map.build_tree(**kwargs),
         cluster_name=cluster_name,
         expected_aio_extensions=expected_extensions,
         expected_aio_custom_locations=expected_custom_locations,
         expected_aio_resources=expected_resources,
         expected_resource_sync_rules=expected_resource_sync_rules,
+        **kwargs,
     )
 
 
@@ -155,27 +162,28 @@ def _assert_tree(
     expected_aio_custom_locations: Optional[List[dict]],
     expected_aio_resources: Optional[List[dict]],
     expected_resource_sync_rules: Optional[List[dict]],
+    category_color: str = "cyan",
 ):
     assert tree.label == f"[green]{cluster_name}"
 
-    assert tree.children[0].label == "[red]extensions"
+    assert tree.children[0].label == f"[{category_color}]extensions"
     if expected_aio_extensions:
         for i in range(len(expected_aio_extensions)):
             tree.children[0].children[i].label == expected_aio_extensions[i]["name"]
 
-    assert tree.children[1].label == "[red]customLocations"
     if expected_aio_custom_locations:
+        assert tree.children[1].label == f"[{category_color}]customLocations"
         for i in range(len(expected_aio_custom_locations)):
             tree.children[1].children[i].label == expected_aio_custom_locations[i]["name"]
 
             if expected_resource_sync_rules:
-                assert tree.children[1].children[i].children[0].label == "[red]resourceSyncRules"
+                assert tree.children[1].children[i].children[0].label == f"[{category_color}]resourceSyncRules"
                 for j in range(len(expected_resource_sync_rules)):
                     tree.children[1].children[i].children[0].children[j].label == expected_resource_sync_rules[i][
                         "name"
                     ]
 
             if expected_aio_resources:
-                assert tree.children[1].children[i].children[1].label == "[red]resources"
+                assert tree.children[1].children[i].children[1].label == f"[{category_color}]resources"
                 for j in range(len(expected_aio_resources)):
                     tree.children[1].children[i].children[1].children[j].label == expected_aio_resources[i]["name"]