Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ boto3
kubernetes
openshift
coldfront >= 1.1.0
pydantic
python-cinderclient # TODO: Set version for OpenStack Clients
python-keystoneclient
python-novaclient
Expand Down
16 changes: 16 additions & 0 deletions src/coldfront_plugin_cloud/attributes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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),
]

###########################################################
Expand Down
65 changes: 65 additions & 0 deletions src/coldfront_plugin_cloud/tests/unit/test_usage_models.py
Original file line number Diff line number Diff line change
@@ -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}})
75 changes: 75 additions & 0 deletions src/coldfront_plugin_cloud/usage_models.py
Original file line number Diff line number Diff line change
@@ -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