unit cov

project-koku · Aug 13, 2024 · c798d65 · c798d65
1 parent 37ca88f
commit c798d65
Showing 1 changed file with 27 additions and 0 deletions.
diff --git a/tests/test_aws_generator.py b/tests/test_aws_generator.py
@@ -326,6 +326,19 @@ def test_update_data(self):
         self.assertEqual(row["product/servicecode"], "AWSDataTransfer")
         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 = {
+            "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):
@@ -390,6 +403,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(