From 392662cc03bedb07b01592f7c6ca4c08ecbf1c9a Mon Sep 17 00:00:00 2001
From: Eva Sebestova <73821679+esebesto@users.noreply.github.com>
Date: Fri, 7 Jun 2024 16:19:12 +0200
Subject: [PATCH] Enable to use AWS EC2 generator instead of Marketplace
 generator (#509)

---
 nise/__init__.py                     |  2 +-
 nise/generators/aws/ec2_generator.py | 11 +++++++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/nise/__init__.py b/nise/__init__.py
index 3d0a53dd..b93d9cc2 100644
--- a/nise/__init__.py
+++ b/nise/__init__.py
@@ -1,3 +1,3 @@
-__version__ = "4.5.3"
+__version__ = "4.5.4"
 
 VERSION = __version__.split(".")
diff --git a/nise/generators/aws/ec2_generator.py b/nise/generators/aws/ec2_generator.py
index 16e0b618..dab94f0b 100644
--- a/nise/generators/aws/ec2_generator.py
+++ b/nise/generators/aws/ec2_generator.py
@@ -121,10 +121,17 @@ def __init__(self, start_date, end_date, currency, payer_account, usage_accounts
     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_description = description.format(cost=cost, inst_type=inst_type)
+        product_name = "Amazon Elastic Compute Cloud"
+        billing_entity = "AWS"
+        if self.attributes:
+            inst_description = self.attributes.get("lineitem_lineitemdescription", inst_description)
+            product_name = self.attributes.get("product_name", product_name)
+            billing_entity = self.attributes.get("billing_entity", billing_entity)
         location, aws_region, avail_zone, _ = self._get_location()
         row = self._add_common_usage_info(row, start, end)
-
+        row["bill/BillingEntity"] = billing_entity
         row["lineItem/ProductCode"] = "AmazonEC2"
         row["lineItem/UsageType"] = f"BoxUsage:{inst_type}"
         row["lineItem/Operation"] = "RunInstances"
@@ -136,7 +143,7 @@ def _update_data(self, row, start, end, **kwargs):
         row["lineItem/BlendedRate"] = rate
         row["lineItem/BlendedCost"] = cost
         row["lineItem/LineItemDescription"] = inst_description
-        row["product/ProductName"] = "Amazon Elastic Compute Cloud"
+        row["product/ProductName"] = product_name
         row["product/clockSpeed"] = "2.8 GHz"
         row["product/currentGeneration"] = "Yes"
         row["product/ecu"] = "14"