From 9a8ad5af5eace7cfa617d2935a802dabb3708d9a Mon Sep 17 00:00:00 2001
From: Luke Couzens <lcouzens@redhat.com>
Date: Mon, 12 Aug 2024 14:11:00 +0100
Subject: [PATCH] Add AWS Savings Plan Negation

---
 nise/__init__.py                              |  2 +-
 .../generators/aws/data_transfer_generator.py | 17 +++++++++--
 nise/generators/aws/ec2_generator.py          | 30 ++++++++++++++++---
 3 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/nise/__init__.py b/nise/__init__.py
index 0eddf277..384e889b 100644
--- a/nise/__init__.py
+++ b/nise/__init__.py
@@ -1,3 +1,3 @@
-__version__ = "4.6.5"
+__version__ = "4.6.6"
 
 VERSION = __version__.split(".")
diff --git a/nise/generators/aws/data_transfer_generator.py b/nise/generators/aws/data_transfer_generator.py
index 3753e454..713fde75 100644
--- a/nise/generators/aws/data_transfer_generator.py
+++ b/nise/generators/aws/data_transfer_generator.py
@@ -45,6 +45,7 @@ def __init__(self, start_date, end_date, currency, payer_account, usage_accounts
         self._rate = float(self.attributes.get("rate", 0)) or None
         self._resource_id = f"i-{self.attributes.get('resource_id', self.fake.ean8())}"
         self._saving = float(self.attributes.get("saving", 0)) or None
+        self._negation = self.attributes.get("negation") or False
         self._tags = self.attributes.get("tags", self._tags)
 
     @property
@@ -81,7 +82,8 @@ def _update_data(self, row, start, end, **kwargs):
 
         resource_id = self._resource_id if self._resource_id else self.fake.ean8()
         rate = self._rate if self._rate else round(uniform(0.12, 0.19), 3)
-        saving = self._saving if self._saving else round(uniform(0.12, 0.19), 3)
+        saving = self._saving
+        negation = self._negation
         amount = self._amount if self._amount else uniform(0.000002, 0.09)
         cost = amount * rate
         trans_desc, operation, description, location1, location2, trans_type, aws_region = self._get_data_transfer(
@@ -118,8 +120,17 @@ def _update_data(self, row, start, end, **kwargs):
         # Overwrite lineItem/LineItemType for items with applied Savings plan
         if saving is not None:
             row["lineItem/LineItemType"] = "SavingsPlanCoveredUsage"
-        self._add_tag_data(row)
-        self._add_category_data(row)
+
+        if negation:
+            row["lineItem/LineItemType"] = "SavingsPlanNegation"
+            row["lineItem/UnblendedCost"] = -abs(cost)
+            row["lineItem/LineItemDescription"] = f"SavingsPlanNegation used by AccountId : {self.payer_account}"
+            row["lineItem/ResourceId"] = None
+            row["lineItem/BlendedCost"] = -abs(cost)
+
+        if not negation:
+            self._add_tag_data(row)
+            self._add_category_data(row)
         return row
 
     def generate_data(self, report_type=None):
diff --git a/nise/generators/aws/ec2_generator.py b/nise/generators/aws/ec2_generator.py
index dab94f0b..74c4313e 100644
--- a/nise/generators/aws/ec2_generator.py
+++ b/nise/generators/aws/ec2_generator.py
@@ -115,12 +115,27 @@ def __init__(self, start_date, end_date, currency, payer_account, usage_accounts
                 instance_type.get("cost"),
                 instance_type.get("rate"),
                 instance_type.get("saving"),
+                instance_type.get("amount", "1"),
+                instance_type.get("negation", False),
                 "${cost} per On Demand Linux {inst_type} Instance Hour",
             )
 
     def _update_data(self, row, start, end, **kwargs):
         """Update data with generator specific data."""
-        inst_type, physical_cores, vcpu, memory, storage, family, cost, rate, saving, description = self._instance_type
+        (
+            inst_type,
+            physical_cores,
+            vcpu,
+            memory,
+            storage,
+            family,
+            cost,
+            rate,
+            saving,
+            amount,
+            negation,
+            description,
+        ) = self._instance_type
 
         inst_description = description.format(cost=cost, inst_type=inst_type)
         product_name = "Amazon Elastic Compute Cloud"
@@ -137,7 +152,7 @@ def _update_data(self, row, start, end, **kwargs):
         row["lineItem/Operation"] = "RunInstances"
         row["lineItem/AvailabilityZone"] = avail_zone
         row["lineItem/ResourceId"] = self._resource_id
-        row["lineItem/UsageAmount"] = "1"
+        row["lineItem/UsageAmount"] = amount
         row["lineItem/UnblendedRate"] = rate
         row["lineItem/UnblendedCost"] = cost
         row["lineItem/BlendedRate"] = rate
@@ -179,9 +194,16 @@ def _update_data(self, row, start, end, **kwargs):
         # Overwrite lineItem/LineItemType for items with applied Savings plan
         if saving is not None:
             row["lineItem/LineItemType"] = "SavingsPlanCoveredUsage"
+        if negation:
+            row["lineItem/LineItemType"] = "SavingsPlanNegation"
+            row["lineItem/UnblendedCost"] = -abs(cost)
+            row["lineItem/LineItemDescription"] = f"SavingsPlanNegation used by AccountId : {self.payer_account}"
+            row["lineItem/ResourceId"] = None
+            row["lineItem/BlendedCost"] = -abs(cost)
 
-        self._add_tag_data(row)
-        self._add_category_data(row)
+        if not negation:
+            self._add_tag_data(row)
+            self._add_category_data(row)
         return row
 
     def generate_data(self, report_type=None):