From 1904e82a8e1041a7c6cc6f143cb2c295a0631840 Mon Sep 17 00:00:00 2001 From: marcoagonzales007 Date: Wed, 5 Aug 2026 15:39:37 -0700 Subject: [PATCH] Move hardcoded quota names from attributes.py to tests --- src/coldfront_plugin_cloud/attributes.py | 28 +-------------- .../commands/calculate_storage_gb_hours.py | 34 ++++++++++--------- .../commands/convert_swift_quota_to_gib.py | 14 ++++---- .../commands/register_default_quotas.py | 26 +++++++------- src/coldfront_plugin_cloud/openstack.py | 17 ++++------ .../tests/attributes.py | 27 +++++++++++++++ .../tests/functional/esi/test_allocations.py | 3 +- .../functional/openshift/test_allocation.py | 3 +- .../openshift_vm/test_allocation.py | 3 +- .../functional/openstack/test_allocation.py | 3 +- .../unit/test_calculate_quota_unit_hours.py | 2 +- .../unit/test_register_default_quotas.py | 2 +- 12 files changed, 84 insertions(+), 78 deletions(-) create mode 100644 src/coldfront_plugin_cloud/tests/attributes.py diff --git a/src/coldfront_plugin_cloud/attributes.py b/src/coldfront_plugin_cloud/attributes.py index 725710a5..334b3d8e 100644 --- a/src/coldfront_plugin_cloud/attributes.py +++ b/src/coldfront_plugin_cloud/attributes.py @@ -87,35 +87,9 @@ class CloudAllocationAttribute: CloudAllocationAttribute(name=ALLOCATION_ALERT, type="Int", is_changeable=True), ] -########################################################### -# OpenStack Quota Attributes -QUOTA_INSTANCES = "OpenStack Compute Instance Quota" -QUOTA_RAM = "OpenStack Compute RAM Quota (MiB)" -QUOTA_VCPU = "OpenStack Compute vCPU Quota" - -QUOTA_VOLUMES = "OpenStack Number of Volumes Quota" -QUOTA_VOLUMES_GB = "OpenStack Volume Quota (GiB)" - -QUOTA_FLOATING_IPS = "OpenStack Floating IP Quota" -QUOTA_NETWORKS = "Openstack Network Quota" - -QUOTA_OBJECT_GB = "OpenStack Swift Quota (GiB)" - +# QUOTE_GPU is still registered because count_gpu_usage command relies on it QUOTA_GPU = "OpenStack GPU Quota" -########################################################### -# OpenShift Quota Attributes -QUOTA_LIMITS_CPU = "OpenShift Limit on CPU Quota" -QUOTA_LIMITS_MEMORY = "OpenShift Limit on RAM Quota (MiB)" -QUOTA_LIMITS_EPHEMERAL_STORAGE_GB = "OpenShift Limit on Ephemeral Storage Quota (GiB)" -QUOTA_REQUESTS_NESE_STORAGE = "OpenShift Request on NESE Storage Quota (GiB)" -QUOTA_REQUESTS_IBM_STORAGE = "OpenShift Request on IBM Storage Quota (GiB)" -QUOTA_REQUESTS_GPU = "OpenShift Request on GPU Quota" -QUOTA_REQUESTS_VM_GPU_A100_SXM4 = "OpenShift Request on GPU A100 SXM4" -QUOTA_REQUESTS_VM_GPU_V100 = "OpenShift Request on GPU V100" -QUOTA_REQUESTS_VM_GPU_H100 = "OpenShift Request on GPU H100" -QUOTA_PVC = "OpenShift Persistent Volume Claims Quota" - ALLOCATION_QUOTA_ATTRIBUTES = [ CloudAllocationAttribute(name=QUOTA_GPU), diff --git a/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py b/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py index e96c4abf..8de00e48 100644 --- a/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py +++ b/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py @@ -322,7 +322,6 @@ def process_invoice_row(allocation, attrs, su_name, rate): f, delimiter=",", quotechar="|", quoting=csv.QUOTE_MINIMAL ) csv_invoice_writer.writerow(InvoiceRow.get_headers()) - for allocation in openstack_allocations: allocation_str = ( f'{allocation.pk} of project "{allocation.project.title}"' @@ -337,6 +336,10 @@ def process_invoice_row(allocation, attrs, su_name, rate): quotaspec.invoice_name, openstack_nese_storage_rate, ) + openshift_storage_rates = { + "OpenShift NESE Storage": openshift_nese_storage_rate, + "OpenShift IBM Scale Storage": openshift_ibm_storage_rate, + } for allocation in openshift_allocations: allocation_str = ( @@ -344,22 +347,21 @@ def process_invoice_row(allocation, attrs, su_name, rate): ) logger.debug(f"Starting billing for allocation {allocation_str}.") - process_invoice_row( - allocation, - [ - attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB, - attributes.QUOTA_REQUESTS_NESE_STORAGE, - ], - "OpenShift NESE Storage", - openshift_nese_storage_rate, - ) + attrs_by_invoice_name = {} + for quota_name, quotaspec in get_storage_quotaspecs(allocation).items(): + attrs_by_invoice_name.setdefault(quotaspec.invoice_name, []).append( + quota_name + ) - process_invoice_row( - allocation, - [attributes.QUOTA_REQUESTS_IBM_STORAGE], - "OpenShift IBM Scale Storage", - openshift_ibm_storage_rate, - ) + for invoice_name, quota_names in attrs_by_invoice_name.items(): + rate = openshift_storage_rates.get(invoice_name) + if rate is None: + logger.warning( + f"No rate configured for invoice name {invoice_name!r} " + f"on allocation {allocation_str}, skipping." + ) + continue + process_invoice_row(allocation, quota_names, invoice_name, rate) if options["upload_to_s3"]: logger.info(f"Uploading to S3 endpoint {options['s3_endpoint_url']}.") diff --git a/src/coldfront_plugin_cloud/management/commands/convert_swift_quota_to_gib.py b/src/coldfront_plugin_cloud/management/commands/convert_swift_quota_to_gib.py index 3d92028c..a236e368 100644 --- a/src/coldfront_plugin_cloud/management/commands/convert_swift_quota_to_gib.py +++ b/src/coldfront_plugin_cloud/management/commands/convert_swift_quota_to_gib.py @@ -9,6 +9,7 @@ from coldfront.core.allocation.models import Allocation, AllocationStatusChoice logging.basicConfig(level=logging.INFO) +SWIFT_QUOTA_ATTRIBUTE = "OpenStack Swift Quota (GiB)" logger = logging.getLogger(__name__) @@ -26,19 +27,20 @@ def handle(self, *args, **options): resources__in=openstack_resources, ) for allocation in openstack_allocations: - if not ( - swift_quota := allocation.get_attribute(attributes.QUOTA_OBJECT_GB) - ): + if not (swift_quota := allocation.get_attribute(SWIFT_QUOTA_ATTRIBUTE)): continue allocation_str = f'{allocation.pk} of project "{allocation.project.title}"' - obj_key = openstack.QUOTA_KEY_MAPPING["object"]["keys"][ - attributes.QUOTA_OBJECT_GB - ] allocator = openstack.OpenStackResourceAllocator( allocation.resources.first(), allocation ) + object_labels = allocator._get_resource_quota_labels_by_service("object") + if not object_labels: + logger.error(f"{allocation_str} has no object quota defined, skipping.") + continue + obj_key = object_labels[0] + project_id = allocation.get_attribute(attributes.ALLOCATION_PROJECT_ID) if not project_id: logger.error(f"{allocation_str} is active but has no Project ID set.") diff --git a/src/coldfront_plugin_cloud/management/commands/register_default_quotas.py b/src/coldfront_plugin_cloud/management/commands/register_default_quotas.py index 05c13010..2dc5ff63 100644 --- a/src/coldfront_plugin_cloud/management/commands/register_default_quotas.py +++ b/src/coldfront_plugin_cloud/management/commands/register_default_quotas.py @@ -37,18 +37,18 @@ def handle(self, *args, **options): # Define quotas for each resource type openshift_quotas = [ { - "display_name": attributes.QUOTA_LIMITS_CPU, + "display_name": "OpenShift Limit on CPU Quota", "quota_label": "limits.cpu", "multiplier": 1, }, { - "display_name": attributes.QUOTA_LIMITS_MEMORY, + "display_name": "OpenShift Limit on RAM Quota (MiB)", "quota_label": "limits.memory", "multiplier": 4096, "unit_suffix": "Mi", }, { - "display_name": attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB, + "display_name": "OpenShift Limit on Ephemeral Storage Quota (GiB)", "quota_label": "limits.ephemeral-storage", "multiplier": 5, "unit_suffix": "Gi", @@ -56,12 +56,12 @@ def handle(self, *args, **options): "invoice_name": OPENSHIFT_STORAGE_INVOICE_NAME, }, { - "display_name": attributes.QUOTA_PVC, + "display_name": "OpenShift Persistent Volume Claims Quota", "quota_label": "persistentvolumeclaims", "multiplier": 2, }, { - "display_name": attributes.QUOTA_REQUESTS_NESE_STORAGE, + "display_name": "OpenShift Request on NESE Storage Quota (GiB)", "quota_label": "ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage", "multiplier": 20, "static_quota": 0, @@ -70,7 +70,7 @@ def handle(self, *args, **options): "invoice_name": OPENSHIFT_STORAGE_INVOICE_NAME, }, { - "display_name": attributes.QUOTA_REQUESTS_GPU, + "display_name": "OpenShift Request on GPU Quota", "quota_label": "requests.nvidia.com/gpu", "multiplier": 0, }, @@ -78,40 +78,40 @@ def handle(self, *args, **options): openstack_quotas = [ { - "display_name": attributes.QUOTA_INSTANCES, + "display_name": "OpenStack Compute Instance Quota", "quota_label": "compute.instances", "multiplier": 1, }, { - "display_name": attributes.QUOTA_VCPU, + "display_name": "OpenStack Compute vCPU Quota", "quota_label": "compute.cores", "multiplier": 1, }, { - "display_name": attributes.QUOTA_RAM, + "display_name": "OpenStack Compute RAM Quota (MiB)", "quota_label": "compute.ram", "multiplier": 4096, }, { - "display_name": attributes.QUOTA_VOLUMES, + "display_name": "OpenStack Number of Volumes Quota", "quota_label": "volume.volumes", "multiplier": 2, }, { - "display_name": attributes.QUOTA_VOLUMES_GB, + "display_name": "OpenStack Volume Quota (GiB)", "quota_label": "volume.gigabytes", "multiplier": 20, "resource_type": STORAGE_RESOURCE_TYPE_NAME, "invoice_name": OPENSTACK_STORAGE_INVOICE_NAME, }, { - "display_name": attributes.QUOTA_FLOATING_IPS, + "display_name": "OpenStack Floating IP Quota", "quota_label": "network.floatingip", "multiplier": 0, "static_quota": 2, }, { - "display_name": attributes.QUOTA_OBJECT_GB, + "display_name": "OpenStack Swift Quota (GiB)", "quota_label": "object.x-account-meta-quota-bytes", "multiplier": 1, "resource_type": STORAGE_RESOURCE_TYPE_NAME, diff --git a/src/coldfront_plugin_cloud/openstack.py b/src/coldfront_plugin_cloud/openstack.py index 43d5935e..70b25e48 100644 --- a/src/coldfront_plugin_cloud/openstack.py +++ b/src/coldfront_plugin_cloud/openstack.py @@ -222,12 +222,12 @@ def _set_object_quota(self, project_id, payload): # Note(knikolla): For consistency with other OpenStack # quotas we're storing this as GB on the attribute and # converting to bytes for Swift. - _, obj_q_mapping = self._extract_quota_label( - self.resource_quotaspecs.root[attributes.QUOTA_OBJECT_GB] - ) - payload[obj_q_mapping] *= GB_IN_BYTES - if payload[obj_q_mapping] <= 0: - payload[obj_q_mapping] = 1 + for obj_q_mapping in self._get_resource_quota_labels_by_service("object"): + if obj_q_mapping not in payload: + continue + payload[obj_q_mapping] *= GB_IN_BYTES + if payload[obj_q_mapping] <= 0: + payload[obj_q_mapping] = 1 self.object(project_id).post_account(headers=payload) except ksa_exceptions.catalog.EndpointNotFound: logger.debug("No swift available, skipping its quota.") @@ -291,10 +291,7 @@ def get_quota(self, project_id): quotas = self._get_network_quota(quotas, project_id) - if object_quotaspec := self.resource_quotaspecs.root.get( - attributes.QUOTA_OBJECT_GB - ): - _, key = self._extract_quota_label(object_quotaspec) + for key in self._get_resource_quota_labels_by_service("object"): try: swift = self.object(project_id).head_account() except ksa_exceptions.catalog.EndpointNotFound: diff --git a/src/coldfront_plugin_cloud/tests/attributes.py b/src/coldfront_plugin_cloud/tests/attributes.py new file mode 100644 index 00000000..e48125d9 --- /dev/null +++ b/src/coldfront_plugin_cloud/tests/attributes.py @@ -0,0 +1,27 @@ +from coldfront_plugin_cloud.attributes import * # noqa: F403 + + +# OpenStack Quota Attributes +QUOTA_INSTANCES = "OpenStack Compute Instance Quota" +QUOTA_RAM = "OpenStack Compute RAM Quota (MiB)" +QUOTA_VCPU = "OpenStack Compute vCPU Quota" + +QUOTA_VOLUMES = "OpenStack Number of Volumes Quota" +QUOTA_VOLUMES_GB = "OpenStack Volume Quota (GiB)" + +QUOTA_FLOATING_IPS = "OpenStack Floating IP Quota" +QUOTA_NETWORKS = "Openstack Network Quota" + +QUOTA_OBJECT_GB = "OpenStack Swift Quota (GiB)" + +# OpenShift Quota Attributes +QUOTA_LIMITS_CPU = "OpenShift Limit on CPU Quota" +QUOTA_LIMITS_MEMORY = "OpenShift Limit on RAM Quota (MiB)" +QUOTA_LIMITS_EPHEMERAL_STORAGE_GB = "OpenShift Limit on Ephemeral Storage Quota (GiB)" +QUOTA_REQUESTS_NESE_STORAGE = "OpenShift Request on NESE Storage Quota (GiB)" +QUOTA_REQUESTS_IBM_STORAGE = "OpenShift Request on IBM Storage Quota (GiB)" +QUOTA_REQUESTS_GPU = "OpenShift Request on GPU Quota" +QUOTA_REQUESTS_VM_GPU_A100_SXM4 = "OpenShift Request on GPU A100 SXM4" +QUOTA_REQUESTS_VM_GPU_V100 = "OpenShift Request on GPU V100" +QUOTA_REQUESTS_VM_GPU_H100 = "OpenShift Request on GPU H100" +QUOTA_PVC = "OpenShift Persistent Volume Claims Quota" diff --git a/src/coldfront_plugin_cloud/tests/functional/esi/test_allocations.py b/src/coldfront_plugin_cloud/tests/functional/esi/test_allocations.py index 77898b8e..6bc15d57 100644 --- a/src/coldfront_plugin_cloud/tests/functional/esi/test_allocations.py +++ b/src/coldfront_plugin_cloud/tests/functional/esi/test_allocations.py @@ -3,7 +3,8 @@ import uuid import time -from coldfront_plugin_cloud import attributes, openstack, esi, tasks +from coldfront_plugin_cloud import openstack, esi, tasks +from coldfront_plugin_cloud.tests import attributes from coldfront_plugin_cloud.tests import base from django.core.management import call_command diff --git a/src/coldfront_plugin_cloud/tests/functional/openshift/test_allocation.py b/src/coldfront_plugin_cloud/tests/functional/openshift/test_allocation.py index 4efb7cfd..53944bda 100644 --- a/src/coldfront_plugin_cloud/tests/functional/openshift/test_allocation.py +++ b/src/coldfront_plugin_cloud/tests/functional/openshift/test_allocation.py @@ -3,7 +3,8 @@ import unittest import uuid -from coldfront_plugin_cloud import attributes, openshift, tasks, utils +from coldfront_plugin_cloud import openshift, tasks, utils +from coldfront_plugin_cloud.tests import attributes from coldfront_plugin_cloud.tests import base from django.core.management import call_command diff --git a/src/coldfront_plugin_cloud/tests/functional/openshift_vm/test_allocation.py b/src/coldfront_plugin_cloud/tests/functional/openshift_vm/test_allocation.py index 97c5323b..24dbbe7f 100644 --- a/src/coldfront_plugin_cloud/tests/functional/openshift_vm/test_allocation.py +++ b/src/coldfront_plugin_cloud/tests/functional/openshift_vm/test_allocation.py @@ -1,7 +1,8 @@ import os import unittest -from coldfront_plugin_cloud import attributes, openshift_vm, tasks +from coldfront_plugin_cloud import openshift_vm, tasks +from coldfront_plugin_cloud.tests import attributes from coldfront_plugin_cloud.tests import base from django.core.management import call_command diff --git a/src/coldfront_plugin_cloud/tests/functional/openstack/test_allocation.py b/src/coldfront_plugin_cloud/tests/functional/openstack/test_allocation.py index a8d7017f..1e3ecff3 100644 --- a/src/coldfront_plugin_cloud/tests/functional/openstack/test_allocation.py +++ b/src/coldfront_plugin_cloud/tests/functional/openstack/test_allocation.py @@ -3,7 +3,8 @@ import uuid import time -from coldfront_plugin_cloud import attributes, openstack, tasks, utils +from coldfront_plugin_cloud import openstack, tasks, utils +from coldfront_plugin_cloud.tests import attributes from coldfront_plugin_cloud.tests import base from django.core.management import call_command diff --git a/src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py b/src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py index bab8a63e..320c4d73 100644 --- a/src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py +++ b/src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py @@ -7,7 +7,7 @@ import freezegun -from coldfront_plugin_cloud import attributes +from coldfront_plugin_cloud.tests import attributes from coldfront_plugin_cloud.tests import base from coldfront_plugin_cloud import utils diff --git a/src/coldfront_plugin_cloud/tests/unit/test_register_default_quotas.py b/src/coldfront_plugin_cloud/tests/unit/test_register_default_quotas.py index d82c8dd5..7398e802 100644 --- a/src/coldfront_plugin_cloud/tests/unit/test_register_default_quotas.py +++ b/src/coldfront_plugin_cloud/tests/unit/test_register_default_quotas.py @@ -3,7 +3,7 @@ from django.core.management import call_command from coldfront_plugin_cloud.tests.base import TestBase -from coldfront_plugin_cloud import attributes +from coldfront_plugin_cloud.tests import attributes class TestRegisterDefaultQuotas(TestBase):