feat: re-enabling billing capture in support bundle

azext_edge/edge/common.py

@@ -145,9 +145,8 @@ class OpsServiceType(ListableEnum):
     orc = "orc"
     akri = "akri"
     deviceregistry = "deviceregistry"
+    billing = "billing"
     dataflow = "dataflow"
-    # TODO: re-enable billing once service is available post 0.6.0 release
-    # billing = "billing"
 
     @classmethod
     def list_check_services(cls):

azext_edge/edge/providers/support/billing.py

@@ -34,8 +34,9 @@ def fetch_pods(
 ):
     # capture billing pods for aio usage
     billing_pods = process_v1_pods(
-        directory_path=ARC_BILLING_DIRECTORY_PATH,
+        directory_path=BILLING_RESOURCE_KIND,
         label_selector=AIO_BILLING_USAGE_NAME_LABEL,
+        prefix_names=["aio-usage"],
         since_seconds=since_seconds,
     )
 
@@ -53,7 +54,7 @@ def fetch_pods(
 
 def fetch_jobs():
     processed = process_jobs(
-        directory_path=ARC_BILLING_DIRECTORY_PATH,
+        directory_path=BILLING_RESOURCE_KIND,
         label_selector=AIO_BILLING_USAGE_NAME_LABEL,
     )
 
@@ -62,7 +63,7 @@ def fetch_jobs():
 
 def fetch_cron_jobs():
     processed = process_cron_jobs(
-        directory_path=ARC_BILLING_DIRECTORY_PATH,
+        directory_path=BILLING_RESOURCE_KIND,
         label_selector=AIO_BILLING_USAGE_NAME_LABEL,
     )
 

azext_edge/edge/providers/support_bundle.py

@@ -46,6 +46,7 @@ def build_bundle(
     from rich.progress import Progress
     from rich.table import Table
 
+    from .support.billing import prepare_bundle as prepare_billing_bundle
     from .support.mq import prepare_bundle as prepare_mq_bundle
     from .support.opcua import prepare_bundle as prepare_opcua_bundle
     from .support.orc import prepare_bundle as prepare_symphony_bundle
@@ -59,9 +60,8 @@ def build_bundle(
     pending_work.pop(OpsServiceType.auto.value)
 
     api_map = {
-        # TODO: re-enable billing once service is available post 0.6.0 release
-        # OpsServiceType.billing.value: {"apis": COMPAT_CLUSTER_CONFIG_APIS, "prepare_bundle": prepare_billing_bundle},
         OpsServiceType.mq.value: {"apis": COMPAT_MQTT_BROKER_APIS, "prepare_bundle": prepare_mq_bundle},
+        OpsServiceType.billing.value: {"apis": COMPAT_CLUSTER_CONFIG_APIS, "prepare_bundle": prepare_billing_bundle},
         OpsServiceType.opcua.value: {
             "apis": COMPAT_OPCUA_APIS,
             "prepare_bundle": prepare_opcua_bundle,

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

@@ -213,7 +213,7 @@ def get_file_map(
         file_map["traces"] = convert_file_names(walk_result[path.join(ops_path, "traces")]["files"])
     elif ops_service == "billing":
         assert len(walk_result) == 2
-        ops_path = path.join(BASE_ZIP_PATH, namespace, "clusterconfig", ops_service)
+        ops_path = path.join(BASE_ZIP_PATH, namespace, ops_service)
         c_path = path.join(BASE_ZIP_PATH, c_namespace, "clusterconfig", ops_service)
         file_map["usage"] = convert_file_names(walk_result[c_path]["files"])
         file_map["__namespaces__"]["usage"] = c_namespace
@@ -262,9 +262,6 @@ def _get_namespace_determinating_files(
         level_1 = walk_result.pop(path.join(BASE_ZIP_PATH, clusterconfig_namespace))
         assert level_1["folders"] == ["clusterconfig"]
         assert not level_1["files"]
-        level_2 = walk_result.pop(path.join(BASE_ZIP_PATH, namespace, "clusterconfig"))
-        assert level_2["folders"] == ["billing"]
-        assert not level_2["files"]
         level_2 = walk_result.pop(path.join(BASE_ZIP_PATH, clusterconfig_namespace, "clusterconfig"))
         assert level_2["folders"] == ["billing"]
         assert not level_2["files"]

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

@@ -69,9 +69,9 @@ def test_create_bundle(init_setup, bundle_dir, mq_traces, ops_service, tracked_f
             walk_result[path.join(BASE_ZIP_PATH, namespace, OpsServiceType.mq.value)]["folders"] = []
 
     # Level 2 and 3 - bottom
-    actual_walk_result = (len(expected_services) + int("clusterconfig" in expected_services))
-
+    actual_walk_result = (len(expected_services) + int(OpsServiceType.billing.value in expected_services))
     assert len(walk_result) == actual_walk_result
+
     for directory in walk_result:
         assert not walk_result[directory]["folders"]
         assert_file_names(walk_result[directory]["files"])
@@ -108,20 +108,13 @@ def _get_expected_services(
     walk_result: Dict[str, Dict[str, List[str]]], ops_service: str , namespace: str
 ) -> List[str]:
     expected_services = [ops_service]
-    # TODO: re-enable billing once service is available post 0.6.0 release
     if ops_service == OpsServiceType.auto.value:
         # these should always be generated
         expected_services = OpsServiceType.list()
         expected_services.remove(OpsServiceType.auto.value)
-        # expected_services.remove(OpsServiceType.billing.value)
         expected_services.append("otel")
-        if walk_result.get(path.join(BASE_ZIP_PATH, namespace, "clusterconfig", "billing")):
-            expected_services.append("clusterconfig")
         # device registry folder will not be created if there are no device registry resources
         if not walk_result.get(path.join(BASE_ZIP_PATH, namespace, OpsServiceType.deviceregistry.value)):
             expected_services.remove(OpsServiceType.deviceregistry.value)
         expected_services.sort()
-    # elif ops_service == OpsServiceType.billing.value:
-    #     expected_services.remove(OpsServiceType.billing.value)
-    #     expected_services.append("clusterconfig")
     return expected_services

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

@@ -4,19 +4,17 @@
 # Licensed under the MIT License. See License file in the project root for license information.
 # ----------------------------------------------------------------------------------------------
 
-import pytest
 from knack.log import get_logger
+from azext_edge.edge.common import OpsServiceType
 from azext_edge.edge.providers.edge_api import CLUSTER_CONFIG_API_V1
 from .helpers import check_custom_resource_files, check_workload_resource_files, get_file_map, run_bundle_command
 
 logger = get_logger(__name__)
 
 
-@pytest.mark.skip(reason="re-enable billing once service is available post 0.6.0 release")
 def test_create_bundle_billing(init_setup, tracked_files):
     """Test for ensuring file names and content. ONLY CHECKS billing."""
-    # TODO: re-enable billing once service is available post 0.6.0 release
-    # ops_service = OpsServiceType.billing.value
+    ops_service = OpsServiceType.billing.value
     ops_service = "billing"
     command = f"az iot ops support create-bundle --ops-service {ops_service}"
     walk_result = run_bundle_command(command=command, tracked_files=tracked_files)

azext_edge/tests/edge/support/test_support_unit.py

@@ -59,6 +59,7 @@
 from azext_edge.edge.providers.support.otel import OTEL_API, OTEL_NAME_LABEL
 from azext_edge.edge.providers.support.shared import NAME_LABEL_FORMAT
 from azext_edge.edge.providers.support_bundle import COMPAT_MQTT_BROKER_APIS
+from azext_edge.tests.edge.support.conftest import add_pod_to_mocked_pods
 
 from ...generators import generate_random_string
 
@@ -75,8 +76,7 @@
         [MQTT_BROKER_API_V1B1, OPCUA_API_V1, DEVICEREGISTRY_API_V1],
         [MQTT_BROKER_API_V1B1, OPCUA_API_V1, ORC_API_V1],
         [MQTT_BROKER_API_V1B1, OPCUA_API_V1, ORC_API_V1, AKRI_API_V0],
-        # TODO: re-enable billing once service is available post 0.6.0 release
-        # [MQ_API_V1B1, OPCUA_API_V1, ORC_API_V1, CLUSTER_CONFIG_API_V1],
+        [MQTT_BROKER_API_V1B1, OPCUA_API_V1, ORC_API_V1, CLUSTER_CONFIG_API_V1],
     ],
     indirect=True,
 )
@@ -114,6 +114,14 @@ def test_create_bundle(
         assert auto_result_no_resources is None
         return
 
+    if CLUSTER_CONFIG_API_V1 in mocked_cluster_resources["param"]:
+        add_pod_to_mocked_pods(
+            mocked_client=mocked_client,
+            expected_pod_map=mocked_list_pods,
+            mock_names=["aio-usage"],
+            mock_init_containers=True,
+        )
+
     since_seconds = random.randint(86400, 172800)
     result = support_bundle(None, bundle_dir=a_bundle_dir, log_age_seconds=since_seconds)
 
@@ -138,6 +146,15 @@ def test_create_bundle(
             )
 
         if api in [CLUSTER_CONFIG_API_V1]:
+            assert_list_pods(
+                mocked_client,
+                mocked_zipfile,
+                mocked_list_pods,
+                label_selector=AIO_BILLING_USAGE_NAME_LABEL,
+                directory_path=BILLING_RESOURCE_KIND,
+                since_seconds=since_seconds,
+                prefix_names=["aio-usage"],
+            )
             assert_list_pods(
                 mocked_client,
                 mocked_zipfile,
@@ -150,7 +167,7 @@ def test_create_bundle(
                 mocked_client,
                 mocked_zipfile,
                 label_selector=AIO_BILLING_USAGE_NAME_LABEL,
-                directory_path=ARC_BILLING_DIRECTORY_PATH,
+                directory_path=BILLING_RESOURCE_KIND,
             )
             assert_list_deployments(
                 mocked_client,
@@ -162,7 +179,7 @@ def test_create_bundle(
                 mocked_client,
                 mocked_zipfile,
                 label_selector=AIO_BILLING_USAGE_NAME_LABEL,
-                directory_path=ARC_BILLING_DIRECTORY_PATH,
+                directory_path=BILLING_RESOURCE_KIND,
             )
             assert_list_replica_sets(
                 mocked_client,
@@ -645,6 +662,9 @@ def assert_list_pods(
                     "namespace: namespace\ntimestamp: '0000-00-00T00:00:00Z'\n",
                 )
 
+            if pod_name not in kwargs.get("prefix_names", []):
+                continue
+
             for container_name in pods_with_container[namespace][pod_name]:
                 data = (
                     f"kind: Pod\nmetadata:\n  name: {pod_name}\n  namespace: {namespace}\nspec:\n  "