From 46e36a607f44b2b5cc5a7fe4d503b8d939022ce3 Mon Sep 17 00:00:00 2001
From: myersCody <cmyers@redhat.com>
Date: Thu, 15 Aug 2024 10:45:43 -0400
Subject: [PATCH] EBS: Add disk_size param to calculate hourly usage & rate

---
 nise/__init__.py                     |  2 +-
 nise/generators/aws/ebs_generator.py | 20 +++++++++-----------
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/nise/__init__.py b/nise/__init__.py
index 4923981a..79e34c47 100644
--- a/nise/__init__.py
+++ b/nise/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "4.6.7"
+__version__ = "4.6.8"
 
 
 VERSION = __version__.split(".")
diff --git a/nise/generators/aws/ebs_generator.py b/nise/generators/aws/ebs_generator.py
index 28a6095f..a83a1cbe 100644
--- a/nise/generators/aws/ebs_generator.py
+++ b/nise/generators/aws/ebs_generator.py
@@ -34,40 +34,38 @@ def __init__(self, start_date, end_date, currency, payer_account, usage_accounts
         """Initialize the EBS generator."""
         super().__init__(start_date, end_date, currency, payer_account, usage_accounts, attributes, tag_cols)
         self._resource_id = "vol-{}".format(self.fake.ean8())
-        self._amount = uniform(0.2, 300.99)
+        self._disk_size = choice([5, 10, 15, 20, 25])
         self._rate = round(uniform(0.02, 0.16), 3)
         self._product_sku = self.fake.pystr(min_chars=12, max_chars=12).upper()
 
         if self.attributes:
             if self.attributes.get("resource_id"):
                 self._resource_id = "vol-{}".format(self.attributes.get("resource_id"))
-            if self.attributes.get("amount"):
-                self._amount = float(self.attributes.get("amount"))
             if self.attributes.get("rate"):
                 self._rate = float(self.attributes.get("rate"))
             if self.attributes.get("product_sku"):
                 self._product_sku = self.attributes.get("product_sku")
             if self.attributes.get("tags"):
                 self._tags = self.attributes.get("tags")
+            if _disk_size := self.attributes.get("disk_size"):
+                self._disk_size = int(_disk_size)
 
     def _get_storage(self):
         """Get storage data."""
         return choice(self.STORAGE)
 
-    def _calculate_cost_per_day(self, start):
-        """EBS charges in a monthly rate.
-        Therefore we need to convert to a daily rate to calculate the cost."""
+    def _calculate_hourly_rate(self, start):
+        """Calculates the houly rate based of the provided monthly rate."""
         num_days_in_month = calendar.monthrange(start.year, start.month)[1]
         hours_in_month = num_days_in_month * 24
-        daily_rate = self._rate / hours_in_month
-        return self._amount * daily_rate
+        return self._rate / hours_in_month
 
     def _update_data(self, row, start, end, **kwargs):
         """Update data with generator specific data."""
         row = self._add_common_usage_info(row, start, end)
-
-        amount = self._amount
-        cost = self._calculate_cost_per_day(start)
+        hourly_rate = self._calculate_hourly_rate(start)
+        cost = round(self._disk_size * hourly_rate, 10)
+        amount = round(cost / self._rate, 10)
         location, aws_region, _, storage_region = self._get_location()
         description = f"${self._rate} per GB-Month of snapshot data stored - {location}"
         burst, max_iops, max_thru, max_vol_size, vol_backed, vol_type = self._get_storage()