Fix parsing

Author: QuanMPhm

src/coldfront_plugin_cloud/management/commands/validate_allocations.py

@@ -90,6 +90,56 @@ def set_default_quota_on_allocation(allocation, allocator, coldfront_attr):
         utils.set_attribute_on_allocation(allocation, coldfront_attr, value)
         return value
 
+    @staticmethod
+    def parse_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 = {
+            "Ki": 2**10,
+            "Mi": 2**20,
+            "Gi": 2**30,
+            "Ti": 2**40,
+            "Pi": 2**50,
+            "Ei": 2**60,
+            "m": 10**-3,
+            "k": 10**3,
+            "K": 10**3,
+            "M": 10**6,
+            "G": 10**9,
+            "T": 10**12,
+            "P": 10**15,
+            "E": 10**18,
+        }
+
+        if quota_str and quota_str != "0":
+            result = re.search(PATTERN, quota_str)
+
+            if result is None:
+                raise CommandError(
+                    f"Unable to parse quota_str = '{quota_str}' for {attr}"
+                )
+
+            value = int(result.groups()[0])
+            unit = result.groups()[1]
+
+            # Convert to number i.e. without any unit suffix
+
+            if unit is not None:
+                quota_str = value * suffix[unit]
+            else:
+                quota_str = value
+
+            # Convert some attributes to units that coldfront uses
+
+            if "RAM" in attr:
+                quota_str = round(quota_str / suffix["Mi"])
+            elif "Storage" in attr:
+                quota_str = round(quota_str / suffix["Gi"])
+        elif quota_str and quota_str == "0":
+            quota_str = 0
+
+        return quota_str
+
     def check_institution_specific_code(self, allocation, apply):
         attr = attributes.ALLOCATION_INSTITUTION_SPECIFIC_CODE
         isc = allocation.get_attribute(attr)
@@ -250,51 +300,7 @@ def handle(self, *args, **options):
 
                 expected_value = allocation.get_attribute(attr)
                 current_value = quota.get(key, None)
-
-                PATTERN = r"([0-9]+)(m|Ki|Mi|Gi|Ti|Pi|Ei|K|M|G|T|P|E)?"
-
-                suffix = {
-                    "Ki": 2**10,
-                    "Mi": 2**20,
-                    "Gi": 2**30,
-                    "Ti": 2**40,
-                    "Pi": 2**50,
-                    "Ei": 2**60,
-                    "m": 10**-3,
-                    "K": 10**3,
-                    "M": 10**6,
-                    "G": 10**9,
-                    "T": 10**12,
-                    "P": 10**15,
-                    "E": 10**18,
-                }
-
-                if current_value and current_value != "0":
-                    result = re.search(PATTERN, current_value)
-
-                    if result is None:
-                        raise CommandError(
-                            f"Unable to parse current_value = '{current_value}' for {attr}"
-                        )
-
-                    value = int(result.groups()[0])
-                    unit = result.groups()[1]
-
-                    # Convert to number i.e. without any unit suffix
-
-                    if unit is not None:
-                        current_value = value * suffix[unit]
-                    else:
-                        current_value = value
-
-                    # Convert some attributes to units that coldfront uses
-
-                    if "RAM" in attr:
-                        current_value = round(current_value / suffix["Mi"])
-                    elif "Storage" in attr:
-                        current_value = round(current_value / suffix["Gi"])
-                elif current_value and current_value == "0":
-                    current_value = 0
+                current_value = self.parse_quota_value(current_value, attr)
 
                 if expected_value is None and current_value is not None:
                     msg = (

src/coldfront_plugin_cloud/tests/unit/test_parse_quota_unit.py

@@ -0,0 +1,29 @@
+from django.core.management.base import CommandError
+
+from coldfront_plugin_cloud.management.commands.validate_allocations import Command
+from coldfront_plugin_cloud.tests import base
+
+
+class TestParseQuotaUnit(base.TestBase):
+    def test_parse_quota_unit(self):
+        parse_quota_unit = Command().parse_quota_value
+        answer_dict = [
+            (("5m", "cpu"), 5 * 10**-3),
+            (("10", "cpu"), 10),
+            (("10k", "cpu"), 10 * 10**3),
+            (("55M", "cpu"), 55 * 10**6),
+            (("2G", "cpu"), 2 * 10**9),
+            (("3T", "cpu"), 3 * 10**12),
+            (("4P", "cpu"), 4 * 10**15),
+            (("5E", "cpu"), 5 * 10**18),
+            (("10", "memory"), 10),
+            (("125Ki", "memory"), 125 * 2**10),
+            (("55Mi", "memory"), 55 * 2**20),
+            (("2Gi", "memory"), 2 * 2**30),
+            (("3Ti", "memory"), 3 * 2**40),
+        ]
+        for (input_value, resource_type), expected in answer_dict:
+            self.assertEqual(parse_quota_unit(input_value, resource_type), expected)
+
+        with self.assertRaises(CommandError):
+            parse_quota_unit("abc", "foo")  # Non-numeric input