diff --git a/requirements.txt b/requirements.txt index 3b441862..b394ba1e 100644 --- a/requirements.txt +++ b/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 diff --git a/src/coldfront_plugin_cloud/attributes.py b/src/coldfront_plugin_cloud/attributes.py index 19d827f8..9bf20a41 100644 --- a/src/coldfront_plugin_cloud/attributes.py +++ b/src/coldfront_plugin_cloud/attributes.py @@ -54,6 +54,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( @@ -65,6 +68,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), ] ########################################################### diff --git a/src/coldfront_plugin_cloud/tests/unit/test_usage_models.py b/src/coldfront_plugin_cloud/tests/unit/test_usage_models.py new file mode 100644 index 00000000..7638a38d --- /dev/null +++ b/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}}) diff --git a/src/coldfront_plugin_cloud/usage_models.py b/src/coldfront_plugin_cloud/usage_models.py new file mode 100644 index 00000000..9d53f0bb --- /dev/null +++ b/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