diff --git a/azext_edge/edge/providers/support/billing.py b/azext_edge/edge/providers/support/billing.py index 17ffad89a..0a266cf69 100644 --- a/azext_edge/edge/providers/support/billing.py +++ b/azext_edge/edge/providers/support/billing.py @@ -57,6 +57,7 @@ def fetch_jobs(): processed = process_jobs( directory_path=BILLING_RESOURCE_KIND, label_selector=AIO_BILLING_USAGE_NAME_LABEL, + prefix_names=[AIO_USAGE_PREFIX], ) return processed diff --git a/azext_edge/tests/edge/support/conftest.py b/azext_edge/tests/edge/support/conftest.py index b6d277968..5b2f0cfdf 100644 --- a/azext_edge/tests/edge/support/conftest.py +++ b/azext_edge/tests/edge/support/conftest.py @@ -6,6 +6,8 @@ from functools import partial from typing import List + +from azext_edge.edge.providers.support.billing import AIO_BILLING_USAGE_NAME_LABEL from ...generators import generate_random_string import pytest @@ -241,8 +243,14 @@ def mocked_list_jobs(mocked_client): from kubernetes.client.models import V1JobList, V1Job, V1ObjectMeta def _handle_list_jobs(*args, **kwargs): - job = V1Job(metadata=V1ObjectMeta(namespace="mock_namespace", name="mock_job")) - job_list = V1JobList(items=[job]) + names = ["mock_job"] + if "label_selector" in kwargs and kwargs["label_selector"] == AIO_BILLING_USAGE_NAME_LABEL: + names.append("aio-usage-job") + + job_list = [] + for name in names: + job_list.append(V1Job(metadata=V1ObjectMeta(namespace="mock_namespace", name=name))) + job_list = V1JobList(items=job_list) return job_list diff --git a/azext_edge/tests/edge/support/test_support_unit.py b/azext_edge/tests/edge/support/test_support_unit.py index a2dbaddcb..553d51915 100644 --- a/azext_edge/tests/edge/support/test_support_unit.py +++ b/azext_edge/tests/edge/support/test_support_unit.py @@ -180,6 +180,7 @@ def test_create_bundle( mocked_zipfile, label_selector=AIO_BILLING_USAGE_NAME_LABEL, directory_path=BILLING_RESOURCE_KIND, + mock_names=["aio-usage-job"], ) assert_list_replica_sets( mocked_client, @@ -599,16 +600,19 @@ def assert_list_jobs( mocked_zipfile, label_selector: str, directory_path: str, + mock_names: Optional[List[str]] = None, ): mocked_client.BatchV1Api().list_job_for_all_namespaces.assert_any_call( label_selector=label_selector, field_selector=None ) - assert_zipfile_write( - mocked_zipfile, - zinfo=f"mock_namespace/{directory_path}/job.mock_job.yaml", - data="kind: Job\nmetadata:\n name: mock_job\n namespace: mock_namespace\n", - ) + mock_names = mock_names or ["mock_job"] + for name in mock_names: + assert_zipfile_write( + mocked_zipfile, + zinfo=f"mock_namespace/{directory_path}/job.{name}.yaml", + data=f"kind: Job\nmetadata:\n name: {name}\n namespace: mock_namespace\n", + ) def assert_list_pods(