1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153
|
from collections import defaultdict
from collections.abc import Iterable
from copy import deepcopy
from datetime import datetime
from typing import Any
from moto.core.base_backend import BackendDict, BaseBackend
from moto.core.common_models import BaseModel
from moto.core.utils import unix_time
from moto.utilities.utils import PARTITION_NAMES
from .exceptions import BudgetMissingLimit, DuplicateRecordException, NotFoundException
class Notification(BaseModel):
def __init__(self, details: dict[str, Any], subscribers: dict[str, Any]):
self.details = details
self.subscribers = subscribers
class Budget(BaseModel):
def __init__(self, budget: dict[str, Any], notifications: list[dict[str, Any]]):
if "BudgetLimit" not in budget and "PlannedBudgetLimits" not in budget:
raise BudgetMissingLimit()
# Storing the budget as a Dict for now - if we need more control, we can always read/write it back
self.budget = budget
self.notifications = [
Notification(details=x["Notification"], subscribers=x["Subscribers"])
for x in notifications
]
self.budget["LastUpdatedTime"] = unix_time()
if "TimePeriod" not in self.budget:
first_day_of_month = datetime.now().replace(
day=1, hour=0, minute=0, second=0, microsecond=0
)
self.budget["TimePeriod"] = {
"Start": unix_time(first_day_of_month),
"End": 3706473600, # "2087-06-15T00:00:00+00:00"
}
def to_dict(self) -> dict[str, Any]:
cp = deepcopy(self.budget)
if "CalculatedSpend" not in cp:
cp["CalculatedSpend"] = {
"ActualSpend": {"Amount": "0", "Unit": "USD"},
"ForecastedSpend": {"Amount": "0", "Unit": "USD"},
}
if self.budget["BudgetType"] == "COST" and "CostTypes" not in cp:
cp["CostTypes"] = {
"IncludeCredit": True,
"IncludeDiscount": True,
"IncludeOtherSubscription": True,
"IncludeRecurring": True,
"IncludeRefund": True,
"IncludeSubscription": True,
"IncludeSupport": True,
"IncludeTax": True,
"IncludeUpfront": True,
"UseAmortized": False,
"UseBlended": False,
}
return cp
def add_notification(
self, details: dict[str, Any], subscribers: dict[str, Any]
) -> None:
self.notifications.append(Notification(details, subscribers))
def delete_notification(self, details: dict[str, Any]) -> None:
self.notifications = [n for n in self.notifications if n.details != details]
def get_notifications(self) -> Iterable[dict[str, Any]]:
return [n.details for n in self.notifications]
class BudgetsBackend(BaseBackend):
"""Implementation of Budgets APIs."""
def __init__(self, region_name: str, account_id: str):
super().__init__(region_name, account_id)
self.budgets: dict[str, dict[str, Budget]] = defaultdict(dict)
def create_budget(
self,
account_id: str,
budget: dict[str, Any],
notifications: list[dict[str, Any]],
) -> None:
budget_name = budget["BudgetName"]
if budget_name in self.budgets[account_id]:
raise DuplicateRecordException(
record_type="budget", record_name=budget_name
)
self.budgets[account_id][budget_name] = Budget(budget, notifications)
def describe_budget(self, account_id: str, budget_name: str) -> dict[str, Any]:
if budget_name not in self.budgets[account_id]:
raise NotFoundException(
f"Unable to get budget: {budget_name} - the budget doesn't exist."
)
return self.budgets[account_id][budget_name].to_dict()
def describe_budgets(self, account_id: str) -> Iterable[dict[str, Any]]:
"""
Pagination is not yet implemented
"""
return [budget.to_dict() for budget in self.budgets[account_id].values()]
def delete_budget(self, account_id: str, budget_name: str) -> None:
if budget_name not in self.budgets[account_id]:
msg = f"Unable to delete budget: {budget_name} - the budget doesn't exist. Try creating it first. "
raise NotFoundException(msg)
self.budgets[account_id].pop(budget_name)
def create_notification(
self,
account_id: str,
budget_name: str,
notification: dict[str, Any],
subscribers: dict[str, Any],
) -> None:
if budget_name not in self.budgets[account_id]:
raise NotFoundException(
"Unable to create notification - the budget doesn't exist."
)
self.budgets[account_id][budget_name].add_notification(
details=notification, subscribers=subscribers
)
def delete_notification(
self, account_id: str, budget_name: str, notification: dict[str, Any]
) -> None:
if budget_name not in self.budgets[account_id]:
raise NotFoundException(
"Unable to delete notification - the budget doesn't exist."
)
self.budgets[account_id][budget_name].delete_notification(details=notification)
def describe_notifications_for_budget(
self, account_id: str, budget_name: str
) -> Iterable[dict[str, Any]]:
"""
Pagination has not yet been implemented
"""
return self.budgets[account_id][budget_name].get_notifications()
budgets_backends = BackendDict(
BudgetsBackend,
"budgets",
use_boto3_regions=False,
additional_regions=PARTITION_NAMES,
)
|