From 2d3edce170664accfcee2a3c16ee82a9b15decd7 Mon Sep 17 00:00:00 2001
From: Luke Couzens <lcouzens@redhat.com>
Date: Tue, 13 Aug 2024 11:44:44 +0100
Subject: [PATCH] unit cov

---
 tests/test_aws_generator.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tests/test_aws_generator.py b/tests/test_aws_generator.py
index 0e955f76..c4051898 100644
--- a/tests/test_aws_generator.py
+++ b/tests/test_aws_generator.py
@@ -327,6 +327,21 @@ def test_update_data(self):
         self.assertEqual(row["product/productFamily"], "Data Transfer")
         self.assertEqual(row[self.cost_category_key], self.cost_category_value)
 
+    def test_update_data_transfer_negation(self):
+        """Test DataTransfer specific update data with negation costs."""
+        self.attributes = {
+            "rate": 20,
+            "amount": 1,
+            "negation": True,
+        }
+        generator = DataTransferGenerator(
+            self.two_hours_ago, self.now, self.currency, self.payer_account, self.usage_accounts, self.attributes
+        )
+        start_row = {}
+        row = generator._update_data(start_row, self.two_hours_ago, self.now)
+
+        self.assertEqual(row["lineItem/LineItemType"], "SavingsPlanNegation")
+
 
 class TestEBSGenerator(AWSGeneratorTestCase):
     """Tests for the EBS Generator type."""
@@ -390,6 +405,20 @@ def test_init_with_attributes(self):
         self.assertEqual(generator._resource_id, "i-" + self.resource_id)
         self.assertEqual(generator._instance_type[:-1], tuple(self.instance_type.values()))
 
+    def test_update_data_ec2_negation(self):
+        """Test EC2 specific update data with negation costs."""
+        self.instance_type = {
+            "negation": True,
+        }
+        self.attributes["instance_type"] = self.instance_type
+        generator = EC2Generator(
+            self.two_hours_ago, self.now, self.currency, self.payer_account, self.usage_accounts, self.attributes
+        )
+        start_row = {}
+        row = generator._update_data(start_row, self.two_hours_ago, self.now)
+
+        self.assertEqual(row["lineItem/LineItemType"], "SavingsPlanNegation")
+
     def test_update_data(self):
         """Test EBS specific update data method."""
         generator = EC2Generator(