Merge pull request #277 from QuanMPhm/270/billing_attr

Add attributes to tracking allocation usage

Author: knikolla

requirements.txt

@@ -3,6 +3,7 @@ boto3
 kubernetes
 openshift
 coldfront >= 1.1.0
+pydantic
 python-cinderclient  # TODO: Set version for OpenStack Clients
 python-keystoneclient
 python-novaclient

src/coldfront_plugin_cloud/attributes.py

@@ -56,6 +56,9 @@ class CloudAllocationAttribute:
 ALLOCATION_PROJECT_ID = "Allocated Project ID"
 ALLOCATION_PROJECT_NAME = "Allocated Project Name"
 ALLOCATION_INSTITUTION_SPECIFIC_CODE = "Institution-Specific Code"
+ALLOCATION_CUMULATIVE_CHARGES = "Cumulative Daily Charges for Month"
+ALLOCATION_PREVIOUS_CHARGES = "Previous Charges"
+ALLOCATION_ALERT = "Monthly Allocation Cost Alert"
 
 ALLOCATION_ATTRIBUTES = [
     CloudAllocationAttribute(
@@ -67,6 +70,19 @@ class CloudAllocationAttribute:
     CloudAllocationAttribute(
         name=ALLOCATION_INSTITUTION_SPECIFIC_CODE, type="Text", is_changeable=True
     ),
+    CloudAllocationAttribute(
+        name=ALLOCATION_CUMULATIVE_CHARGES,
+        type="Text",
+        is_private=True,
+        is_changeable=True,
+    ),
+    CloudAllocationAttribute(
+        name=ALLOCATION_PREVIOUS_CHARGES,
+        type="Text",
+        is_private=True,
+        is_changeable=True,
+    ),
+    CloudAllocationAttribute(name=ALLOCATION_ALERT, type="Int", is_changeable=True),
 ]
 
 ###########################################################

src/coldfront_plugin_cloud/tests/unit/test_usage_models.py

@@ -0,0 +1,65 @@
+from decimal import Decimal
+from pydantic import ValidationError
+
+from coldfront_plugin_cloud import usage_models
+from coldfront_plugin_cloud.tests import base
+
+
+class TestUsageModels(base.TestBase):
+    def test_usage_info(self):
+        # valid: values coerced to Decimal
+        ui = usage_models.UsageInfo(
+            root={"su-a": Decimal("1.5"), "su-b": 2, "su-c": "3.25"}
+        )
+        self.assertIsInstance(ui.root, dict)
+        self.assertEqual(ui.root["su-a"], Decimal("1.5"))
+        self.assertEqual(ui.root["su-b"], Decimal("2"))
+        self.assertEqual(ui.root["su-c"], Decimal("3.25"))
+
+        # invalid: non-numeric string should raise ValidationError
+        with self.assertRaises(ValidationError):
+            usage_models.UsageInfo(root={"su-x": "not-a-number"})
+
+    def test_daily_charges_dict(self):
+        # Valid CumulativeChargesDict with YYYY-MM-DD keys
+        data = {
+            "2025-11-29": {"su1": Decimal("1.0")},
+            "2025-11-30": {"su1": Decimal("3.5"), "su2": Decimal("2.0")},
+        }
+        daily = usage_models.CumulativeChargesDict(root=data)
+        # total_charges sums across all dates and SUs
+        self.assertEqual(daily.total_charges, Decimal("5.5"))
+
+        # Empty dict -> totals should be zero/empty
+        empty = usage_models.CumulativeChargesDict(root={})
+        self.assertEqual(empty.total_charges, Decimal("0.0"))
+
+        # Invalid date key format should raise ValidationError
+        with self.assertRaises(ValidationError):
+            usage_models.CumulativeChargesDict(root={"2025-13-01": {"su": 1.0}})
+
+        with self.assertRaises(ValidationError):
+            usage_models.CumulativeChargesDict(root={"2025-01": {"su": 1.0}})
+
+        # Different months should raise ValidationError
+        with self.assertRaises(ValidationError):
+            usage_models.CumulativeChargesDict(
+                root={"2025-12-01": {"su": 1.0}, "2026-12-01": {"su": 1.0}}
+            )
+
+    def test_previous_charges_dict(self):
+        # Monthly (PreviousChargesDict) requires YYYY-MM keys
+        prev_data = {
+            "2025-11": {"suA": Decimal("5.0")},
+            "2025-12": {"suA": Decimal("2.5"), "suB": Decimal("1.0")},
+        }
+        prev = usage_models.PreviousChargesDict(root=prev_data)
+        self.assertEqual(
+            prev.total_charges_by_su,
+            {"suA": Decimal("7.5"), "suB": Decimal("1.0")},
+        )
+        self.assertEqual(prev.total_charges, Decimal("8.5"))
+
+        # Invalid month format should raise ValidationError
+        with self.assertRaises(ValidationError):
+            usage_models.PreviousChargesDict(root={"2025-11-01": {"su": 1.0}})

src/coldfront_plugin_cloud/usage_models.py

@@ -0,0 +1,75 @@
+import datetime
+import functools
+from decimal import Decimal
+from typing import Annotated, TypeVar
+
+import pydantic
+
+
+def validate_date_str(v: str) -> str:
+    datetime.datetime.strptime(v, "%Y-%m-%d")
+    return v
+
+
+def validate_month_str(v: str) -> str:
+    datetime.datetime.strptime(v, "%Y-%m")
+    return v
+
+
+DateField = Annotated[str, pydantic.AfterValidator(validate_date_str)]
+MonthField = Annotated[str, pydantic.AfterValidator(validate_month_str)]
+
+
+class UsageInfo(pydantic.RootModel[dict[str, Decimal]]):
+    pass
+
+
+T = TypeVar("T", bound=str)
+
+
+class ChargesDict(pydantic.RootModel[dict[T, UsageInfo]]):
+    @functools.cached_property
+    def most_recent_date(self) -> DateField:
+        """Leverage lexical ordering of YYYY-MM-DD and YYYY-MM strings."""
+        return max(self.root.keys()) if self.root else ""
+
+
+class CumulativeChargesDict(ChargesDict[DateField]):
+    @pydantic.model_validator(mode="after")
+    def check_month(self):
+        # Ensure all keys are in the same month
+        if self.root:
+            months = set()
+            for date_str in self.root.keys():
+                months.add(
+                    datetime.datetime.strptime(date_str, "%Y-%m-%d").strftime("%Y-%m")
+                )
+
+            if len(months) != 1:
+                raise ValueError("All dates must be within the same month")
+        return self
+
+    @functools.cached_property
+    def total_charges(self) -> Decimal:
+        total = Decimal("0.00")
+        if most_recent_charges := self.root.get(self.most_recent_date):
+            for su_charge in most_recent_charges.root.values():
+                total += su_charge
+        return total
+
+
+class PreviousChargesDict(ChargesDict[MonthField]):
+    @functools.cached_property
+    def total_charges_by_su(self) -> dict[str, Decimal]:
+        total = {}
+        for usage_info in self.root.values():
+            for su_name, charge in usage_info.root.items():
+                total[su_name] = total.get(su_name, Decimal("0.00")) + charge
+        return total
+
+    @functools.cached_property
+    def total_charges(self) -> Decimal:
+        total = Decimal("0.00")
+        for su_charge in self.total_charges_by_su.values():
+            total += su_charge
+        return total