Merge pull request #202 from QuanMPhm/feature/add_attr

Implemented feature to add new attributes to existing Openshift allocations

Author: knikolla

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

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",
+            },
+        )