[COST-4367] Allow legal entity setting for AWS and fix date issue

nise/__init__.py

@@ -1,3 +1,3 @@
-__version__ = "4.4.11"
+__version__ = "4.4.12"
 
 VERSION = __version__.split(".")

nise/__main__.py

@@ -742,7 +742,7 @@ def calculate_start_date(start_date):
     elif start_date == "today":
         generated_start_date = today().replace(hour=0, minute=0, second=0)
     elif start_date and isinstance(start_date, datetime.date):
-        generated_start_date = start_date
+        generated_start_date = datetime.datetime(start_date.year, start_date.month, start_date.day)
     elif start_date:
         generated_start_date = date_parser.parse(start_date)
     else:
@@ -760,7 +760,7 @@ def calculate_end_date(start_date, end_date):
         elif end_date == "today":
             generated_end_date = today().replace(hour=0, minute=0, second=0)
         elif end_date and isinstance(end_date, datetime.date):
-            generated_end_date = end_date
+            generated_end_date = datetime.datetime(end_date.year, end_date.month, end_date.day)
         else:
             generated_end_date = date_parser.parse(end_date)
     except TypeError:

nise/generators/aws/aws_generator.py

@@ -333,14 +333,20 @@ def _get_location(self):
             location = choice(REGIONS)
         return location
 
+    def _get_legal_entity(self):
+        """Pick legal entity."""
+        if self.attributes and self.attributes.get("legal_entity"):
+            return self.attributes.get("legal_entity")
+        return "Amazon Web Services, Inc."
+
     def _add_common_usage_info(self, row, start, end, **kwargs):
         """Add common usage information."""
         row["lineItem/UsageAccountId"] = choice(self.usage_accounts)
         row["lineItem/LineItemType"] = "Usage"
         row["lineItem/UsageStartDate"] = start
         row["lineItem/UsageEndDate"] = end
         row["lineItem/CurrencyCode"] = self.currency
-        row["lineItem/LegalEntity"] = "Amazon Web Services, Inc."
+        row["lineItem/LegalEntity"] = self._get_legal_entity()
         return row
 
     def _add_tag_data(self, row):

tests/test_aws_generator.py

@@ -156,6 +156,21 @@ def test_get_location(self):
         location = generator._get_location()
         self.assertIn("us-west-1", location)
 
+    def test_get_legal_entity(self):
+        """Test the _get_legal_entity method."""
+        two_hours_ago = (self.now - self.one_hour) - self.one_hour
+        generator = TestGenerator(two_hours_ago, self.now, self.currency, self.payer_account, self.usage_accounts)
+        legal_entity = generator._get_legal_entity()
+        self.assertEqual(legal_entity, "Amazon Web Services, Inc.")
+
+        attributes = {}
+        attributes["legal_entity"] = "Corey"
+        generator = TestGenerator(
+            two_hours_ago, self.now, self.currency, self.payer_account, self.usage_accounts, attributes
+        )
+        legal_entity = generator._get_legal_entity()
+        self.assertEqual(legal_entity, "Corey")
+
 
 class AWSGeneratorTestCase(TestCase):
     """Test Base for specific generator classes."""