[COST-4744] Azure Data Transfer Generator (#498)

* [COST-4744] Add Azure Data Transfer Generator to nise

nise/__init__.py

@@ -1,3 +1,3 @@
-__version__ = "4.4.16"
+__version__ = "4.4.17"
 
 VERSION = __version__.split(".")

nise/generators/azure/__init__.py

@@ -20,6 +20,7 @@
 from nise.generators.azure.azure_generator import AzureGenerator  # noqa: F401
 from nise.generators.azure.bandwidth_generator import BandwidthGenerator  # noqa: F401
 from nise.generators.azure.ccsp_generator import CCSPGenerator  # noqa: F401
+from nise.generators.azure.data_transfer_generator import DTGenerator  # noqa: F401
 from nise.generators.azure.sql_database_generator import SQLGenerator  # noqa: F401
 from nise.generators.azure.storage_generator import StorageGenerator  # noqa: F401
 from nise.generators.azure.virtual_machine_generator import VMGenerator  # noqa: F401

nise/generators/azure/azure_generator.py

@@ -187,6 +187,7 @@ def __init__(self, start_date, end_date, currency, account_info, attributes=None
         self._meter_cache = {}
         self._billing_currency = currency
         self._additional_info = None
+        self._data_direction = None
         # Version 2 fields
         self._invoice_section_id = None
         self._invoice_section_name = None
@@ -225,7 +226,7 @@ def _get_resource_info(self, meter_id, service_meter, ex_resource, service_info)
         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 = self._get_additional_info()
+        additional_info = self._get_additional_info(meter_name)
         if self._instance_id:
             self._consumed, second_part = accts_str = self._get_accts_str(self._service_name)
             self._resource_type = self._consumed + "/" + second_part
@@ -292,7 +293,7 @@ def _get_location(self):
             location = choice(self.RESOURCE_LOCATION)
         return location
 
-    def _get_additional_info(self):
+    def _get_additional_info(self, meter_name=None):
         """Pick additional info."""
         if self._additional_info:
             return self._additional_info

nise/generators/azure/data_transfer_generator.py

@@ -0,0 +1,81 @@
+#
+# Copyright 2024 Red Hat, Inc.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+"""Module for azure bandwidth data generation."""
+from random import choice
+
+from nise.generators.azure.azure_generator import AzureGenerator
+
+
+class DTGenerator(AzureGenerator):
+    """Generator for Virtual Machine data."""
+
+    ACCTS_STR = {
+        "Virtual Network": ("microsoft.compute", "publicIPAddresses"),
+    }
+    SERVICE_METER = {
+        "in": (
+            "Virtual Network Private Link",
+            "Virtual Network Private Link",
+            "Standard Data Processed - Ingress",
+            "1 GB",
+        ),
+        "out": (
+            "Virtual Network Private Link",
+            "Virtual Network Private Link",
+            "Standard Data Processed - Egress",
+            "1 GB",
+        ),
+    }
+    SERVICE_FAMILIES = "Networking"
+    EXAMPLE_RESOURCE = (
+        ("RG1", "mysa1"),
+        ("RG1", "costmgmtacct1234"),
+        ("RG2", "mysa1"),
+        ("RG2", "costmgmtacct1234"),
+        ("costmgmt", "mysa1"),
+        ("costmgmt", "costmgmtacct1234"),
+        ("hccm", "mysa1"),
+        ("hccm", "costmgmtacct1234"),
+    )
+    ADDITIONAL_INFO = {
+        "Standard Data Processed - Egress": {
+            "DataTransferDirection": "DataTrOut",
+        },
+        "Standard Data Processed - Ingress": {
+            "DataTransferDirection": "DataTrIn",
+        },
+    }
+
+    def __init__(self, start_date, end_date, currency, account_info, attributes=None):
+        """Initialize the data transfer generator."""
+        self._service_name = "Virtual Network"
+        super().__init__(start_date, end_date, currency, account_info, attributes)
+
+    def _get_additional_info(self, meter_name):
+        """Pick additional info."""
+        return self.ADDITIONAL_INFO.get(meter_name, {})
+
+    def _get_cached_meter_values(self, meter_id, service_meter):
+        """Return meter cached meter data to ensure meter_id and values are consistent."""
+        if not self._meter_cache.get(f"{meter_id}_{self._data_direction}"):
+            if self._data_direction:
+                self._meter_cache[f"{meter_id}_{self._data_direction}"] = service_meter.get(self._data_direction)
+            else:
+                self._meter_cache[f"{meter_id}_{self._data_direction}"] = service_meter.get(
+                    choice(list(service_meter))
+                )
+        return self._meter_cache.get(f"{meter_id}_{self._data_direction}")

nise/report.py

@@ -53,6 +53,7 @@
 from nise.generators.aws import VPCGenerator
 from nise.generators.azure import BandwidthGenerator
 from nise.generators.azure import CCSPGenerator
+from nise.generators.azure import DTGenerator
 from nise.generators.azure import SQLGenerator
 from nise.generators.azure import StorageGenerator
 from nise.generators.azure import VMGenerator
@@ -749,6 +750,7 @@ def azure_create_report(options):  # noqa: C901
             {"generator": StorageGenerator, "attributes": {}},
             {"generator": VMGenerator, "attributes": {}},
             {"generator": VNGenerator, "attributes": {}},
+            {"generator": DTGenerator, "attributes": {}},
         ]
         accounts_list = None
 

tests/test_azure_generator.py

@@ -24,6 +24,7 @@
 from nise.generators.azure import AZURE_COLUMNS_V2
 from nise.generators.azure import AzureGenerator
 from nise.generators.azure import BandwidthGenerator
+from nise.generators.azure import DTGenerator
 from nise.generators.azure import SQLGenerator
 from nise.generators.azure import StorageGenerator
 from nise.generators.azure import VMGenerator
@@ -441,3 +442,75 @@ def test_update_data_with_attributes(self):
             else:
                 self.assertEqual(row["SubscriptionId"], self.payer_account)
                 self.assertEqual(row["ResourceId"], self.instance_id)
+
+
+class TestDTGenerator(AzureGeneratorTestCase):
+    """Tests for the VM Generator type."""
+
+    ADDITIONAL_INFO_KEYS = {"Standard Data Processed - Egress", "Standard Data Processed - Ingress"}
+
+    def test_init_no_attributes(self):
+        """Test the init wihout attributes."""
+        generator = DTGenerator(
+            self.two_hours_ago, self.now, self.currency, self.account_info, attributes={"empty": "dictionary"}
+        )
+        self.assertEqual(generator._service_name, "Virtual Network")
+
+    def test_init_with_attributes(self):
+        """Test the unique init options for VM."""
+        default_generators = [
+            DTGenerator(self.two_hours_ago, self.now, self.currency, self.account_info, self.attributes),
+            DTGenerator(self.two_hours_ago, self.now, self.currency, self.account_info, self.attributes_v2),
+        ]
+        for generator in default_generators:
+            self.assertEqual(generator._service_name, "Virtual Network")
+            self.assertEqual(generator._meter_id, self.meter_id)
+            self.assertEqual(generator._tags, self.tags)
+            self.assertEqual(generator._usage_quantity, self.usage_quantity)
+            self.assertEqual(generator._resource_rate, self.resource_rate)
+            self.assertEqual(generator._pre_tax_cost, self.pre_tax_cost)
+
+    def test_update_data(self):
+        """Test that row is updated."""
+        generator = DTGenerator(self.two_hours_ago, self.now, self.currency, self.account_info)
+        start_row = {}
+        row = generator._update_data(start_row, self.two_hours_ago, self.now)
+        self.assertEqual(row["ConsumedService"], "microsoft.compute")
+        self.assertEqual(row["ResourceType"], "microsoft.compute/publicIPAddresses")
+
+    def test_update_data_with_attributes(self):
+        """Test that row is updated."""
+        directional_attributes = {"data_direction": "in"}
+        directional_attributes.update(self.attributes)
+        directional_attributes_v2 = {"data_direction": "in"}
+        directional_attributes_v2.update(self.attributes_v2)
+        default_generators = [
+            DTGenerator(
+                self.two_hours_ago,
+                self.now,
+                self.currency,
+                self.account_info,
+                directional_attributes,
+            ),
+            DTGenerator(
+                self.two_hours_ago,
+                self.now,
+                self.currency,
+                self.account_info,
+                directional_attributes_v2,
+            ),
+        ]
+        for generator in default_generators:
+            start_row = {}
+            row = generator._update_data(start_row, self.two_hours_ago, self.now)
+            self.assertEqual(row["Tags"], json.dumps(self.tags))
+            if generator.azure_columns == AZURE_COLUMNS:
+                self.assertEqual(row["SubscriptionGuid"], self.payer_account)
+                self.assertEqual(row["InstanceId"], self.instance_id)
+                self.assertEqual(row["MeterSubcategory"], "Virtual Network Private Link")
+                self.assertEqual(row["MeterName"], "Standard Data Processed - Ingress")
+            else:
+                self.assertEqual(row["SubscriptionId"], self.payer_account)
+                self.assertEqual(row["ResourceId"], self.instance_id)
+                self.assertEqual(row["MeterSubCategory"], "Virtual Network Private Link")
+                self.assertEqual(row["MeterName"], "Standard Data Processed - Ingress")