diff --git a/src/coldfront_plugin_cloud/management/commands/validate_allocations.py b/src/coldfront_plugin_cloud/management/commands/validate_allocations.py index 4c6e46e..71b9bf8 100644 --- a/src/coldfront_plugin_cloud/management/commands/validate_allocations.py +++ b/src/coldfront_plugin_cloud/management/commands/validate_allocations.py @@ -82,6 +82,14 @@ def sync_openshift_project_labels(project_id, allocator, apply): f"Labels updated for Openshift project {project_id}: {', '.join(missing_or_incorrect_labels)}" ) + @staticmethod + def set_default_quota_on_allocation(allocation, allocator, coldfront_attr): + uqm = tasks.UNIT_QUOTA_MULTIPLIERS[allocator.resource_type] + value = allocation.quantity * uqm.get(coldfront_attr, 0) + value += tasks.STATIC_QUOTA[allocator.resource_type].get(coldfront_attr, 0) + utils.set_attribute_on_allocation(allocation, coldfront_attr, value) + return value + def check_institution_specific_code(self, allocation, apply): attr = attributes.ALLOCATION_INSTITUTION_SPECIFIC_CODE isc = allocation.get_attribute(attr) @@ -299,19 +307,41 @@ def handle(self, *args, **options): ) msg = f"{msg} Attribute set to match current quota." logger.warning(msg) - elif 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) + 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"]: - try: - allocator.set_quota(project_id) + if options["apply"]: + expected_value = self.set_default_quota_on_allocation( + allocation, allocator, attr + ) logger.warning( - f"Quota for allocation {project_id} was out of date. Reapplied!" + f"Added default quota for {attr} to allocation {allocation_str} to {expected_value}" ) - except Exception as e: - logger.error(f"setting openshift quota failed: {e}") - continue + + 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 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 ad66d15..b4e7fa0 100644 --- a/src/coldfront_plugin_cloud/tests/functional/openshift/test_allocation.py +++ b/src/coldfront_plugin_cloud/tests/functional/openshift/test_allocation.py @@ -1,6 +1,7 @@ import os import time import unittest +from unittest import mock import uuid from coldfront_plugin_cloud import attributes, openshift, tasks, utils @@ -324,3 +325,54 @@ def test_create_incomplete(self): user.username, user.username ) ) + + @mock.patch.object( + tasks, + "UNIT_QUOTA_MULTIPLIERS", + { + "openshift": { + attributes.QUOTA_LIMITS_CPU: 1, + } + }, + ) + def test_allocation_new_attribute(self): + """When a new attribute is introduced, but pre-existing allocations don't have it""" + user = self.new_user() + project = self.new_project(pi=user) + allocation = self.new_allocation(project, self.resource, 2) + allocator = openshift.OpenShiftResourceAllocator(self.resource, allocation) + + tasks.activate_allocation(allocation.pk) + allocation.refresh_from_db() + + project_id = allocation.get_attribute(attributes.ALLOCATION_PROJECT_ID) + + self.assertEqual(allocation.get_attribute(attributes.QUOTA_LIMITS_CPU), 2 * 1) + + quota = allocator.get_quota(project_id) + self.assertEqual( + quota, + { + "limits.cpu": "2", + }, + ) + + # Add a new attribute for Openshift + tasks.UNIT_QUOTA_MULTIPLIERS["openshift"][attributes.QUOTA_LIMITS_MEMORY] = 4096 + + call_command("validate_allocations", apply=True) + allocation.refresh_from_db() + + self.assertEqual(allocation.get_attribute(attributes.QUOTA_LIMITS_CPU), 2 * 1) + self.assertEqual( + allocation.get_attribute(attributes.QUOTA_LIMITS_MEMORY), 2 * 4096 + ) + + quota = allocator.get_quota(project_id) + self.assertEqual( + quota, + { + "limits.cpu": "2", + "limits.memory": "8Gi", + }, + )