File: models.py

package info (click to toggle)
python-moto 5.1.18-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 116,520 kB
  • sloc: python: 636,725; javascript: 181; makefile: 39; sh: 3
file content (153 lines) | stat: -rw-r--r-- 5,797 bytes parent folder | download
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,
)