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..e39cf926 100644 --- a/nise/generators/aws/aws_generator.py +++ b/nise/generators/aws/aws_generator.py @@ -333,6 +333,12 @@ 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) @@ -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/nise/generators/azure/azure_generator.py b/nise/generators/azure/azure_generator.py index 38ace7fc..a48234ab 100644 --- a/nise/generators/azure/azure_generator.py +++ b/nise/generators/azure/azure_generator.py @@ -190,6 +190,7 @@ def __init__(self, start_date, end_date, currency, account_info, attributes=None self._resource_type = None self._meter_cache = {} self._billing_currency = currency + self._additional_info = None # Version 2 fields self._invoice_section_id = None self._invoice_section_name = None @@ -223,12 +224,12 @@ def _get_cached_meter_values(self, meter_id, service_meter): self._meter_cache[meter_id] = choice(service_meter) return self._meter_cache.get(meter_id) - def _get_resource_info(self, meter_id, service_meter, ex_resource, add_info, service_info): + def _get_resource_info(self, meter_id, service_meter, ex_resource, service_info): """Return resource information.""" service_tier, meter_sub, meter_name, units_of_measure = self._get_cached_meter_values(meter_id, service_meter) service_info_2 = choice(service_info) resource_group, resource_name = choice(ex_resource) - additional_info = choice(add_info) + additional_info = self._get_additional_info() if self._instance_id: self._consumed, second_part = accts_str = self._get_accts_str(self._service_name) self._resource_type = self._consumed + "/" + second_part @@ -295,6 +296,13 @@ def _get_location(self): location = choice(self.RESOURCE_LOCATION) return location + def _get_additional_info(self): + """Pick additional info.""" + if self._additional_info: + return self._additional_info + else: + return choice(self.ADDITIONAL_INFO) + def _add_common_usage_info(self, row, start, end, **kwargs): """Add common usage information.""" if self.azure_columns == AZURE_COLUMNS_V2: @@ -341,9 +349,7 @@ def _update_data(self, row, start, end, **kwargs): units_of_measure, additional_info, service_info_2, - ) = self._get_resource_info( - meter_id, self.SERVICE_METER, self.EXAMPLE_RESOURCE, self.ADDITIONAL_INFO, self.SERVICE_INFO_2 - ) + ) = self._get_resource_info(meter_id, self.SERVICE_METER, self.EXAMPLE_RESOURCE, self.SERVICE_INFO_2) if not additional_info: additional_info = "" if not service_info_2: diff --git a/tests/test_aws_generator.py b/tests/test_aws_generator.py index d1c34e70..470f2459 100644 --- a/tests/test_aws_generator.py +++ b/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.""" diff --git a/tests/test_azure_generator.py b/tests/test_azure_generator.py index 6b41a298..8c0025d0 100644 --- a/tests/test_azure_generator.py +++ b/tests/test_azure_generator.py @@ -172,6 +172,19 @@ def test_get_location(self): location = generator._get_location() self.assertIn("US East", location) + def test_get_additional_info(self): + """Test the _get_additional_info method.""" + two_hours_ago = (self.now - self.one_hour) - self.one_hour + generator = TestGenerator(two_hours_ago, self.now, self.currency, self.account_info) + add_info = generator._get_additional_info() + self.assertIsNone(add_info) + + attributes = {} + attributes["additional_info"] = {"VCPU": "1"} + generator = TestGenerator(two_hours_ago, self.now, self.currency, self.account_info, attributes) + add_info = generator._get_additional_info() + self.assertIn("VCPU", add_info) + class AzureGeneratorTestCase(TestCase): """Test Base for specific generator classes."""