feat: support bundle support multiple --svc (#394)

azext_edge/edge/_help.py

@@ -89,6 +89,14 @@ def load_iotops_help():
         - name: Include secretstore resources in the support bundle.
           text: >
             az iot ops support create-bundle --ops-service secretstore
+
+        - name: Include multiple services in the support bundle with single --ops-service flag.
+          text: >
+            az iot ops support create-bundle --ops-service broker opcua deviceregistry
+
+        - name: Include multiple services in the support bundle with multiple --ops-service flags.
+          text: >
+            az iot ops support create-bundle --ops-service broker --ops-service opcua --ops-service deviceregistry
     """
 
     helps[

azext_edge/edge/commands_edge.py

@@ -11,7 +11,6 @@
 from azure.cli.core.azclierror import ArgumentUsageError
 from knack.log import get_logger
 
-from .common import OpsServiceType
 from .providers.base import DEFAULT_NAMESPACE, load_config_context
 from .providers.check.common import ResourceOutputDetailLevel
 from .providers.edge_api import META_API_V1B1
@@ -30,17 +29,17 @@
 def support_bundle(
     cmd,
     log_age_seconds: int = 60 * 60 * 24,
-    ops_service: str = OpsServiceType.auto.value,
     bundle_dir: Optional[str] = None,
     include_mq_traces: Optional[bool] = None,
     context_name: Optional[str] = None,
+    ops_services: Optional[List[str]] = None,
 ) -> Union[Dict[str, Any], None]:
     load_config_context(context_name=context_name)
     from .providers.support_bundle import build_bundle
 
     bundle_path: PurePath = get_bundle_path(bundle_dir=bundle_dir)
     return build_bundle(
-        ops_service=ops_service,
+        ops_services=ops_services,
         bundle_path=str(bundle_path),
         log_age_seconds=log_age_seconds,
         include_mq_traces=include_mq_traces,

azext_edge/edge/common.py

@@ -145,7 +145,6 @@ class OpsServiceType(ListableEnum):
     IoT Operations service type.
     """
 
-    auto = "auto"
     mq = "broker"
     opcua = "opcua"
     akri = "akri"

azext_edge/edge/params.py

@@ -130,11 +130,14 @@ def load_iotops_arguments(self, _):
 
     with self.argument_context("iot ops support") as context:
         context.argument(
-            "ops_service",
+            "ops_services",
+            nargs="+",
+            action="extend",
             options_list=["--ops-service", "--svc"],
             choices=CaseInsensitiveList(OpsServiceType.list()),
             help="The IoT Operations service the support bundle creation should apply to. "
-            "If auto is selected, the operation will detect which services are available.",
+            "If no service is provided, the operation will default to capture all services. "
+            "--ops-service can be used one or more times.",
         )
         context.argument(
             "log_age_seconds",
@@ -580,14 +583,12 @@ def load_iotops_arguments(self, _):
             help="Asset endpoint profile name.",
         )
         context.argument(
-            "instance_name",
-            options_list=["--instance"],
-            help="Instance name to associate the created asset with."
+            "instance_name", options_list=["--instance"], help="Instance name to associate the created asset with."
         )
         context.argument(
             "instance_resource_group",
             options_list=["--instance-resource-group", "--ig"],
-            help="Instance resource group. If not provided, asset resource group will be used."
+            help="Instance resource group. If not provided, asset resource group will be used.",
         )
         context.argument(
             "instance_subscription",
@@ -964,7 +965,7 @@ def load_iotops_arguments(self, _):
         context.argument(
             "instance_resource_group",
             options_list=["--instance-resource-group", "--ig"],
-            help="Instance resource group. If not provided, asset endpoint profile resource group will be used."
+            help="Instance resource group. If not provided, asset endpoint profile resource group will be used.",
         )
         context.argument(
             "instance_subscription",
@@ -994,7 +995,7 @@ def load_iotops_arguments(self, _):
             options_list=["--authentication-mode", "--am"],
             help="Authentication Mode.",
             arg_group="Authentication",
-            arg_type=get_enum_type(AEPAuthModes)
+            arg_type=get_enum_type(AEPAuthModes),
         )
         context.argument(
             "certificate_reference",

azext_edge/edge/providers/support_bundle.py

@@ -40,9 +40,9 @@
 
 
 def build_bundle(
-    ops_service: str,
     bundle_path: str,
     log_age_seconds: Optional[int] = None,
+    ops_services: Optional[List[str]] = None,
     include_mq_traces: Optional[bool] = None,
 ):
     from rich.live import Live
@@ -77,7 +77,6 @@ def collect_default_works(
         pending_work["meta"] = prepare_meta_bundle(log_age_seconds, deployed_meta_apis)
 
     pending_work = {k: {} for k in OpsServiceType.list()}
-    pending_work.pop(OpsServiceType.auto.value)
 
     api_map = {
         OpsServiceType.mq.value: {"apis": COMPAT_MQTT_BROKER_APIS, "prepare_bundle": prepare_mq_bundle},
@@ -112,35 +111,43 @@ def collect_default_works(
         },
     }
 
-    for service_moniker, api_info in api_map.items():
-        if ops_service in [OpsServiceType.auto.value, service_moniker]:
-            deployed_apis = api_info["apis"].get_deployed() if api_info["apis"] else None
-
-            if not deployed_apis and service_moniker not in [
-                OpsServiceType.schemaregistry.value,
-                OpsServiceType.akri.value,
-            ]:
-                expected_api_version = api_info["apis"].as_str()
-                logger.warning(
-                    f"The following API(s) were not detected {expected_api_version}. "
-                    f"CR capture for {service_moniker} will be skipped. "
-                    "Still attempting capture of runtime resources..."
-                )
-
-            # still try fetching other resources even crds are not available due to api version mismatch
-            bundle_method = api_info["prepare_bundle"]
-            # Check if the function takes a second argument
-            # TODO: Change to kwargs based pattern
-            if service_moniker == OpsServiceType.deviceregistry.value:
-                bundle = bundle_method(deployed_apis)
-            elif service_moniker == OpsServiceType.mq.value:
-                bundle = bundle_method(log_age_seconds, deployed_apis, include_mq_traces)
-            elif service_moniker in [OpsServiceType.schemaregistry.value, OpsServiceType.akri.value]:
-                bundle = bundle_method(log_age_seconds)
-            else:
-                bundle = bundle_method(log_age_seconds, deployed_apis)
-
-            pending_work[service_moniker].update(bundle)
+    if not ops_services:
+        parsed_ops_services = OpsServiceType.list()
+    else:
+        # remove duplicates
+        parsed_ops_services = list(set(ops_services))
+
+    for ops_service in parsed_ops_services:
+        # assign key and value to service_moniker and api_info
+        service_moniker = [k for k, _ in api_map.items() if k == ops_service][0]
+        api_info = api_map.get(service_moniker)
+        deployed_apis = api_info["apis"].get_deployed() if api_info["apis"] else None
+
+        if not deployed_apis and service_moniker not in [
+            OpsServiceType.schemaregistry.value,
+            OpsServiceType.akri.value,
+        ]:
+            expected_api_version = api_info["apis"].as_str()
+            logger.warning(
+                f"The following API(s) were not detected {expected_api_version}. "
+                f"CR capture for {service_moniker} will be skipped. "
+                "Still attempting capture of runtime resources..."
+            )
+
+        # still try fetching other resources even crds are not available due to api version mismatch
+        bundle_method = api_info["prepare_bundle"]
+        # Check if the function takes a second argument
+        # TODO: Change to kwargs based pattern
+        if service_moniker == OpsServiceType.deviceregistry.value:
+            bundle = bundle_method(deployed_apis)
+        elif service_moniker == OpsServiceType.mq.value:
+            bundle = bundle_method(log_age_seconds, deployed_apis, include_mq_traces)
+        elif service_moniker in [OpsServiceType.schemaregistry.value, OpsServiceType.akri.value]:
+            bundle = bundle_method(log_age_seconds)
+        else:
+            bundle = bundle_method(log_age_seconds, deployed_apis)
+
+        pending_work[service_moniker].update(bundle)
 
     collect_default_works(pending_work, log_age_seconds)
 

azext_edge/tests/edge/support/create_bundle_int/test_auto_int.py

@@ -25,6 +25,9 @@ def generate_bundle_test_cases() -> List[Tuple[str, bool, Optional[str]]]:
     # case = ops_service, mq_traces, bundle_dir
     cases = [(service, False, "support_bundles") for service in OpsServiceType.list()]
     cases.append((OpsServiceType.mq.value, True, None))
+
+    # test "all services" bundle
+    cases.append((None, False, None))
     return cases
 
 
@@ -36,20 +39,23 @@ def test_create_bundle(init_setup, bundle_dir, mq_traces, ops_service, tracked_f
     if ops_service == OpsServiceType.arccontainerstorage.value:
         pytest.skip("arccontainerstorage is not generated in aio namespace")
 
-    command = f"az iot ops support create-bundle --broker-traces {mq_traces} " + "--ops-service {0}"
+    command = f"az iot ops support create-bundle --broker-traces {mq_traces} "
     if bundle_dir:
-        command += f" --bundle-dir {bundle_dir}"
+        command += f" --bundle-dir {bundle_dir} "
         try:
             mkdir(bundle_dir)
             tracked_files.append(bundle_dir)
         except FileExistsError:
             pass
-    walk_result, _ = run_bundle_command(command=command.format(ops_service), tracked_files=tracked_files)
+
     # generate second bundle as close as possible
-    if ops_service != OpsServiceType.auto.value:
-        auto_walk_result, _ = run_bundle_command(
-            command=command.format(OpsServiceType.auto.value), tracked_files=tracked_files
+    if ops_service:
+        walk_result, _ = run_bundle_command(
+            command=command + f"--ops-service {ops_service}", tracked_files=tracked_files
         )
+        auto_walk_result, _ = run_bundle_command(command=command, tracked_files=tracked_files)
+    else:
+        walk_result, _ = run_bundle_command(command=command, tracked_files=tracked_files)
 
     # Level 0 - top
     namespaces = process_top_levels(walk_result, ops_service)
@@ -64,7 +70,7 @@ def test_create_bundle(init_setup, bundle_dir, mq_traces, ops_service, tracked_f
     assert not level_1["files"]
 
     # Check and take out mq traces:
-    if mq_traces and ops_service in [OpsServiceType.auto.value, OpsServiceType.mq.value]:
+    if mq_traces and ops_service == OpsServiceType.mq.value:
         mq_level = walk_result.pop(path.join(BASE_ZIP_PATH, aio_namespace, OpsServiceType.mq.value, "traces"), {})
         if mq_level:
             assert not mq_level["folders"]
@@ -94,7 +100,7 @@ def test_create_bundle(init_setup, bundle_dir, mq_traces, ops_service, tracked_f
         assert_file_names(walk_result[directory]["files"])
 
     # check service is within auto
-    if ops_service != OpsServiceType.auto.value:
+    if ops_service:
         expected_folders = [[]]
         if mq_traces and ops_service == OpsServiceType.mq.value:
             expected_folders.append(["traces"])
@@ -115,12 +121,7 @@ def test_create_bundle(init_setup, bundle_dir, mq_traces, ops_service, tracked_f
 def _get_expected_services(
     walk_result: Dict[str, Dict[str, List[str]]], ops_service: str, namespace: str
 ) -> List[str]:
-    expected_services = [ops_service]
-    if ops_service == OpsServiceType.auto.value:
-        # these should always be generated
-        expected_services = OpsServiceType.list()
-        expected_services.remove(OpsServiceType.auto.value)
-        expected_services.sort()
+    expected_services = [ops_service] if ops_service else OpsServiceType.list()
 
     # device registry folder will not be created if there are no device registry resources
     if (

azext_edge/tests/edge/support/create_bundle_int/test_meta_int.py

@@ -5,30 +5,20 @@
 # ----------------------------------------------------------------------------------------------
 
 from knack.log import get_logger
-from azext_edge.edge.common import OpsServiceType
 from azext_edge.edge.providers.edge_api import META_API_V1B1
-from .helpers import (
-    check_custom_resource_files,
-    check_workload_resource_files,
-    get_file_map,
-    run_bundle_command
-)
+from .helpers import check_custom_resource_files, check_workload_resource_files, get_file_map, run_bundle_command
 
 logger = get_logger(__name__)
 
 
 def test_create_bundle_meta(init_setup, tracked_files):
     """Test for ensuring file names and content. ONLY CHECKS meta."""
     # dir for unpacked files
-    ops_service = OpsServiceType.auto.value
-    command = f"az iot ops support create-bundle --ops-service {ops_service}"
+    command = "az iot ops support create-bundle"
     walk_result, bundle_path = run_bundle_command(command=command, tracked_files=tracked_files)
     file_map = get_file_map(walk_result, "meta")["aio"]
 
-    check_custom_resource_files(
-        file_objs=file_map,
-        resource_api=META_API_V1B1
-    )
+    check_custom_resource_files(file_objs=file_map, resource_api=META_API_V1B1)
 
     expected_workload_types = ["deployment", "pod", "replicaset", "service"]
     optional_workload_types = ["job"]
@@ -37,10 +27,6 @@ def test_create_bundle_meta(init_setup, tracked_files):
     check_workload_resource_files(
         file_objs=file_map,
         expected_workload_types=expected_workload_types,
-        prefixes=[
-            "aio-operator",
-            "aio-pre-install-job",
-            "aio-post-install-job"
-        ],
-        bundle_path=bundle_path
+        prefixes=["aio-operator", "aio-pre-install-job", "aio-post-install-job"],
+        bundle_path=bundle_path,
     )

azext_edge/tests/edge/support/test_akri_support_unit.py

@@ -40,7 +40,7 @@ def test_create_bundle_akri(
     since_seconds = random.randint(86400, 172800)
     result = support_bundle(
         None,
-        ops_service=OpsServiceType.akri.value,
+        ops_services=[OpsServiceType.akri.value],
         bundle_dir=a_bundle_dir,
         log_age_seconds=since_seconds,
     )

azext_edge/tests/edge/support/test_secretstore_support_unit.py

@@ -42,7 +42,7 @@ def test_create_bundle_ssc(
     since_seconds = random.randint(86400, 172800)
     result = support_bundle(
         None,
-        ops_service=OpsServiceType.secretstore.value,
+        ops_services=[OpsServiceType.secretstore.value],
         bundle_dir=a_bundle_dir,
         log_age_seconds=since_seconds,
     )

azext_edge/tests/edge/support/test_support_unit.py

@@ -428,7 +428,7 @@ def test_create_bundle_crd_work(
     mocked_get_config_map: Mock,
     mocked_assemble_crd_work,
 ):
-    support_bundle(None, ops_service=OpsServiceType.mq.value, bundle_dir=a_bundle_dir)
+    support_bundle(None, ops_services=[OpsServiceType.mq.value], bundle_dir=a_bundle_dir)
 
     if mocked_cluster_resources["param"] == []:
         mocked_root_logger.warning.assert_called_with(
@@ -877,7 +877,9 @@ def test_create_bundle_mq_traces(
     mocked_mq_get_traces,
     mocked_get_config_map,
 ):
-    result = support_bundle(None, ops_service=OpsServiceType.mq.value, bundle_dir=a_bundle_dir, include_mq_traces=True)
+    result = support_bundle(
+        None, ops_services=[OpsServiceType.mq.value], bundle_dir=a_bundle_dir, include_mq_traces=True
+    )
 
     assert result["bundlePath"]
     mocked_mq_get_traces.assert_called_once()
@@ -918,7 +920,7 @@ def test_create_bundle_arc_agents(
     since_seconds = random.randint(86400, 172800)
     result = support_bundle(
         None,
-        ops_service=OpsServiceType.deviceregistry.value,
+        ops_services=[OpsServiceType.deviceregistry.value],
         bundle_dir=a_bundle_dir,
         log_age_seconds=since_seconds,
     )
@@ -976,7 +978,7 @@ def test_create_bundle_schemas(
     since_seconds = random.randint(86400, 172800)
     result = support_bundle(
         None,
-        ops_service=OpsServiceType.schemaregistry.value,
+        ops_services=[OpsServiceType.schemaregistry.value],
         bundle_dir=a_bundle_dir,
         log_age_seconds=since_seconds,
     )