From 68f645b4d1061363d0d6952b47a679bf397fa0c8 Mon Sep 17 00:00:00 2001
From: Corey Goodfred <cgoodfre@redhat.com>
Date: Wed, 6 Dec 2023 13:08:20 -0500
Subject: [PATCH 1/2] allow legal entity setting for AWS and fix date issue

---
 nise/__init__.py                     |  2 +-
 nise/__main__.py                     |  4 ++--
 nise/generators/aws/aws_generator.py |  8 +++++++-
 tests/test_aws_generator.py          | 15 +++++++++++++++
 4 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/nise/__init__.py b/nise/__init__.py
index dcfc3b6e..e3c44d1d 100644
--- a/nise/__init__.py
+++ b/nise/__init__.py
@@ -1,3 +1,3 @@
-__version__ = "4.4.10"
+__version__ = "4.4.12"
 
 VERSION = __version__.split(".")
diff --git a/nise/__main__.py b/nise/__main__.py
index 4db5ac22..d9cbb3fb 100644
--- a/nise/__main__.py
+++ b/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:
diff --git a/nise/generators/aws/aws_generator.py b/nise/generators/aws/aws_generator.py
index e8dd61ee..1da25462 100644
--- a/nise/generators/aws/aws_generator.py
+++ b/nise/generators/aws/aws_generator.py
@@ -332,6 +332,12 @@ def _get_location(self):
         else:
             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."""
@@ -340,7 +346,7 @@ def _add_common_usage_info(self, row, start, end, **kwargs):
         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):
diff --git a/tests/test_aws_generator.py b/tests/test_aws_generator.py
index d1c34e70..5c5a282f 100644
--- a/tests/test_aws_generator.py
+++ b/tests/test_aws_generator.py
@@ -155,6 +155,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):

From 2591f1c7108fbda53107c09c94f07ed67cc4547c Mon Sep 17 00:00:00 2001
From: Corey Goodfred <cgoodfre@redhat.com>
Date: Wed, 6 Dec 2023 13:12:09 -0500
Subject: [PATCH 2/2] pre-commit

---
 nise/generators/aws/aws_generator.py | 2 +-
 tests/test_aws_generator.py          | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/nise/generators/aws/aws_generator.py b/nise/generators/aws/aws_generator.py
index 1da25462..e39cf926 100644
--- a/nise/generators/aws/aws_generator.py
+++ b/nise/generators/aws/aws_generator.py
@@ -332,7 +332,7 @@ def _get_location(self):
         else:
             location = choice(REGIONS)
         return location
-    
+
     def _get_legal_entity(self):
         """Pick legal entity."""
         if self.attributes and self.attributes.get("legal_entity"):
diff --git a/tests/test_aws_generator.py b/tests/test_aws_generator.py
index 5c5a282f..470f2459 100644
--- a/tests/test_aws_generator.py
+++ b/tests/test_aws_generator.py
@@ -155,7 +155,7 @@ 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