[COST-4306] Allow setting of additional info for Azure

nise/__init__.py

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

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:

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."""