Refactored validate_allocations to be more resource-agnostic

Some validation steps still require different code paths for
OpenStack and OpenShift allocations. Particularly in quota validation

Quota validation will now behave the same for both resource types.
OpenStack's particular use of default quotas is reflected in a new
test in `openstack/test_allocation.py`

OpenStack integration code is slightly changed to
better handle missing object storage quotas

Author: QuanMPhm

src/coldfront_plugin_cloud/management/commands/validate_allocations.py

@@ -23,6 +23,9 @@
 class Command(BaseCommand):
     help = "Validates quotas and users in resource allocations."
 
+    OPENSTACK_RESOURCE_NAMES = ["OpenStack", "ESI"]
+    OPENSHIFT_RESOURCE_NAMES = ["OpenShift", "Openshift Virtualization"]
+
     def add_arguments(self, parser):
         parser.add_argument(
             "--apply",
@@ -91,7 +94,7 @@ def set_default_quota_on_allocation(allocation, allocator, coldfront_attr):
         return value
 
     @staticmethod
-    def parse_quota_value(quota_str: str | None, attr: str) -> int | None:
+    def parse_openshift_quota_value(quota_str: str | None, attr: str) -> int | None:
         PATTERN = r"([0-9]+)(m|k|Ki|Mi|Gi|Ti|Pi|Ei|K|M|G|T|P|E)?"
 
         suffix = {
@@ -140,6 +143,97 @@ def parse_quota_value(quota_str: str | None, attr: str) -> int | None:
 
         return quota_str
 
+    def validate_project_exists(self, allocator, project_id, resource_name):
+        if resource_name in self.OPENSHIFT_RESOURCE_NAMES:
+            allocator._get_project(project_id)
+        elif resource_name in self.OPENSTACK_RESOURCE_NAMES:
+            allocator.identity.projects.get(project_id)
+
+    def validate_quotas(
+        self,
+        allocator,
+        project_id,
+        allocation,
+        allocation_str,
+        resource_name,
+        apply: bool,
+    ):
+        quota = allocator.get_quota(project_id)
+        for attr in tasks.get_expected_attributes(allocator):
+            # quota_key = Command.get_quota_key(attr, resource_name)
+            # Get quota key
+            if resource_name in self.OPENSHIFT_RESOURCE_NAMES:
+                key_with_lambda = allocator.QUOTA_KEY_MAPPING.get(attr, None)
+                # This gives me just the plain key str
+                quota_key = list(key_with_lambda(1).keys())[0]
+            elif resource_name in self.OPENSTACK_RESOURCE_NAMES:
+                quota_key = allocator.QUOTA_KEY_MAPPING_ALL_KEYS.get(attr, None)
+                if not quota_key:
+                    # Note(knikolla): Some attributes are only maintained
+                    # for bookkeeping purposes and do not have a
+                    # corresponding quota set on the service.
+                    continue
+
+            expected_value = allocation.get_attribute(attr)
+            current_value = quota.get(quota_key, None)
+            # expected_value, current_value = Command.parse_quota_values(expected_value, current_value, attr, resource_name)
+
+            # parse quota values
+            if resource_name in self.OPENSHIFT_RESOURCE_NAMES:
+                current_value = Command.parse_openshift_quota_value(current_value, attr)
+            elif resource_name in self.OPENSTACK_RESOURCE_NAMES:
+                obj_key = openstack.OpenStackResourceAllocator.QUOTA_KEY_MAPPING[
+                    "object"
+                ]["keys"][attributes.QUOTA_OBJECT_GB]
+                if quota_key == obj_key and expected_value <= 0:
+                    expected_value = 1
+                    current_value = int(
+                        allocator.object(project_id).head_account().get(obj_key)
+                    )
+
+            if current_value is None and expected_value is None:
+                msg = (
+                    f"Value for quota for {attr} is not set anywhere"
+                    f" on allocation {allocation_str}"
+                )
+
+                if apply:
+                    expected_value = Command.set_default_quota_on_allocation(
+                        allocation, allocator, attr
+                    )
+                    msg = f"Added default quota for {attr} to allocation {allocation_str} to {expected_value}"
+                logger.warning(msg)
+            elif current_value is not None and expected_value is None:
+                msg = (
+                    f'Attribute "{attr}" expected on allocation {allocation_str} but not set.'
+                    f" Current quota is {current_value}."
+                )
+
+                if apply:
+                    utils.set_attribute_on_allocation(allocation, attr, current_value)
+                    expected_value = (
+                        current_value  # To pass `current_value != expected_value` check
+                    )
+                    msg = f"{msg} Attribute set to match current quota."
+                logger.warning(msg)
+
+            if current_value != expected_value:
+                msg = (
+                    f"Value for quota for {attr} = {current_value} does not match expected"
+                    f" value of {expected_value} on allocation {allocation_str}"
+                )
+                logger.warning(msg)
+
+                if apply:
+                    try:
+                        allocator.set_quota(project_id)
+                        logger.warning(
+                            f"Quota for allocation {project_id} was out of date. Reapplied!"
+                        )
+                    except Exception as e:
+                        logger.error(f"setting openshift quota failed: {e}")
+                        continue
+
     def check_institution_specific_code(self, allocation, apply):
         attr = attributes.ALLOCATION_INSTITUTION_SPECIFIC_CODE
         isc = allocation.get_attribute(attr)
@@ -152,202 +246,55 @@ def check_institution_specific_code(self, allocation, apply):
                 logger.warn(f'Attribute "{attr}" added to allocation {alloc_str}')
 
     def handle(self, *args, **options):
-        # Deal with Openstack and ESI resources
-        openstack_resources = Resource.objects.filter(
-            resource_type__name__in=["OpenStack", "ESI"]
-        )
-        openstack_allocations = Allocation.objects.filter(
-            resources__in=openstack_resources,
-            status=AllocationStatusChoice.objects.get(name="Active"),
-        )
-        for allocation in openstack_allocations:
-            self.check_institution_specific_code(allocation, options["apply"])
-            allocation_str = f'{allocation.pk} of project "{allocation.project.title}"'
-            msg = f"Starting resource validation for allocation {allocation_str}."
-            logger.debug(msg)
-
-            failed_validation = False
-
-            allocator = tasks.find_allocator(allocation)
-
-            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.")
-                continue
-
-            try:
-                allocator.identity.projects.get(project_id)
-            except http.NotFound:
-                logger.error(
-                    f"{allocation_str} has Project ID {project_id}. But"
-                    f" no project found in OpenStack."
-                )
-                continue
-
-            quota = allocator.get_quota(project_id)
-
-            failed_validation = Command.sync_users(
-                project_id, allocation, allocator, options["apply"]
+        for resource_name in (
+            self.OPENSTACK_RESOURCE_NAMES + self.OPENSHIFT_RESOURCE_NAMES
+        ):
+            resource = Resource.objects.filter(resource_type__name=resource_name)
+            allocations = Allocation.objects.filter(
+                resources__in=resource,
+                status=AllocationStatusChoice.objects.get(name="Active"),
             )
 
-            obj_key = openstack.OpenStackResourceAllocator.QUOTA_KEY_MAPPING["object"][
-                "keys"
-            ][attributes.QUOTA_OBJECT_GB]
+            for allocation in allocations:
+                allocation_str = (
+                    f'{allocation.pk} of project "{allocation.project.title}"'
+                )
+                logger.debug(
+                    f"Starting resource validation for allocation {allocation_str}."
+                )
 
-            for attr in tasks.get_expected_attributes(allocator):
-                key = allocator.QUOTA_KEY_MAPPING_ALL_KEYS.get(attr, None)
-                if not key:
-                    # Note(knikolla): Some attributes are only maintained
-                    # for bookkeeping purposes and do not have a
-                    # corresponding quota set on the service.
-                    continue
+                allocator = tasks.find_allocator(allocation)
+                project_id = allocation.get_attribute(attributes.ALLOCATION_PROJECT_ID)
 
-                expected_value = allocation.get_attribute(attr)
-                current_value = quota.get(key, None)
-                if key == obj_key and expected_value <= 0:
-                    expected_obj_value = 1
-                    current_value = int(
-                        allocator.object(project_id).head_account().get(obj_key)
-                    )
-                    if current_value != expected_obj_value:
-                        failed_validation = True
-                        msg = (
-                            f"Value for quota for {attr} = {current_value} does not match expected"
-                            f" value of {expected_obj_value} on allocation {allocation_str}"
-                        )
-                        logger.warning(msg)
-                elif expected_value is None and current_value:
-                    msg = (
-                        f'Attribute "{attr}" expected on allocation {allocation_str} but not set.'
-                        f" Current quota is {current_value}."
-                    )
-                    if options["apply"]:
-                        utils.set_attribute_on_allocation(
-                            allocation, attr, current_value
-                        )
-                        msg = f"{msg} Attribute set to match current quota."
-                    logger.warning(msg)
-                elif not current_value == expected_value:
-                    failed_validation = True
-                    msg = (
-                        f"Value for quota for {attr} = {current_value} does not match expected"
-                        f" value of {expected_value} on allocation {allocation_str}"
+                # Check project ID is set
+                if not project_id:
+                    logger.error(
+                        f"{allocation_str} is active but has no Project ID set."
                     )
-                    logger.warning(msg)
+                    continue
 
-            if failed_validation and options["apply"]:
+                # Check project exists in remote cluster
                 try:
-                    allocator.set_quota(
-                        allocation.get_attribute(attributes.ALLOCATION_PROJECT_ID)
-                    )
-                except Exception as e:
+                    self.validate_project_exists(allocator, project_id, resource_name)
+                except http.NotFound:
                     logger.error(
-                        f"setting {allocation.resources.first()} quota failed: {e}"
+                        f"{allocation_str} has Project ID {project_id}. But"
+                        f" no project found in {resource_name}."
                     )
                     continue
-                logger.warning(
-                    f"Quota for allocation {allocation_str} was out of date. Reapplied!"
-                )
-
-        # Deal with OpenShift and Openshift VM
-
-        openshift_resources = Resource.objects.filter(
-            resource_type__name__in=["OpenShift", "Openshift Virtualization"]
-        )
-        openshift_allocations = Allocation.objects.filter(
-            resources__in=openshift_resources,
-            status=AllocationStatusChoice.objects.get(name="Active"),
-        )
-
-        for allocation in openshift_allocations:
-            self.check_institution_specific_code(allocation, options["apply"])
-            allocation_str = f'{allocation.pk} of project "{allocation.project.title}"'
-            logger.debug(
-                f"Starting resource validation for allocation {allocation_str}."
-            )
 
-            allocator = tasks.find_allocator(allocation)
-
-            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.")
-                continue
-
-            try:
-                allocator._get_project(project_id)
-            except http.NotFound:
-                logger.error(
-                    f"{allocation_str} has Project ID {project_id}. But"
-                    f" no project found in OpenShift."
-                )
-                continue
-
-            quota = allocator.get_quota(project_id)
-
-            failed_validation = Command.sync_users(
-                project_id, allocation, allocator, options["apply"]
-            )
-            Command.sync_openshift_project_labels(
-                project_id, allocator, options["apply"]
-            )
-
-            for attr in tasks.get_expected_attributes(allocator):
-                key_with_lambda = allocator.QUOTA_KEY_MAPPING.get(attr, None)
-
-                # This gives me just the plain key
-                key = list(key_with_lambda(1).keys())[0]
-
-                expected_value = allocation.get_attribute(attr)
-                current_value = quota.get(key, None)
-                current_value = self.parse_quota_value(current_value, attr)
-
-                if expected_value is None and current_value is not None:
-                    msg = (
-                        f'Attribute "{attr}" expected on allocation {allocation_str} but not set.'
-                        f" Current quota is {current_value}."
+                # Check institution code, users, labels, and quotas
+                self.check_institution_specific_code(allocation, options["apply"])
+                Command.sync_users(project_id, allocation, allocator, options["apply"])
+                if resource_name in self.OPENSHIFT_RESOURCE_NAMES:
+                    Command.sync_openshift_project_labels(
+                        project_id, allocator, options["apply"]
                     )
-                    if options["apply"]:
-                        utils.set_attribute_on_allocation(
-                            allocation, attr, current_value
-                        )
-                        msg = f"{msg} Attribute set to match current quota."
-                    logger.warning(msg)
-                else:
-                    # We just checked the case where the quota value is set in the cluster
-                    # but not in coldfront. This is the only case the cluster value is the
-                    # "source of truth" for the quota value
-                    # If the coldfront value is set, it is always the source of truth.
-                    # But first, we need to check if the quota value is set anywhere at all.
-                    # TODO (Quan): Refactor these if statements so that we can remove this comment block
-                    if current_value is None and expected_value is None:
-                        msg = (
-                            f"Value for quota for {attr} is not set anywhere"
-                            f" on allocation {allocation_str}"
-                        )
-                        logger.warning(msg)
-
-                        if options["apply"]:
-                            expected_value = self.set_default_quota_on_allocation(
-                                allocation, allocator, attr
-                            )
-                            logger.warning(
-                                f"Added default quota for {attr} to allocation {allocation_str} to {expected_value}"
-                            )
-
-                    if not (current_value == expected_value):
-                        msg = (
-                            f"Value for quota for {attr} = {current_value} does not match expected"
-                            f" value of {expected_value} on allocation {allocation_str}"
-                        )
-                        logger.warning(msg)
-
-                        if options["apply"]:
-                            try:
-                                allocator.set_quota(project_id)
-                                logger.warning(
-                                    f"Quota for allocation {project_id} was out of date. Reapplied!"
-                                )
-                            except Exception as e:
-                                logger.error(f"setting openshift quota failed: {e}")
-                                continue
+                self.validate_quotas(
+                    allocator,
+                    project_id,
+                    allocation,
+                    allocation_str,
+                    resource_name,
+                    options["apply"],
+                )

src/coldfront_plugin_cloud/openstack.py

@@ -263,16 +263,17 @@ def get_quota(self, project_id):
         key = self.QUOTA_KEY_MAPPING["object"]["keys"][attributes.QUOTA_OBJECT_GB]
         try:
             swift = self.object(project_id).head_account()
-            quotas[key] = int(int(swift.get(key)) / GB_IN_BYTES)
         except ksa_exceptions.catalog.EndpointNotFound:
             logger.debug("No swift available, skipping its quota.")
         except swiftclient.exceptions.ClientException as e:
             if e.http_status == 403:
                 self._init_rgw_for_project(project_id)
-                swift = self.object(project_id).head_account()
-                quotas[key] = int(int(swift.get(key)) / GB_IN_BYTES)
             else:
                 raise
+
+        try:
+            swift = self.object(project_id).head_account()
+            quotas[key] = int(int(swift.get(key)) / GB_IN_BYTES)
         except (ValueError, TypeError):
             logger.info("No swift quota set.")