From db45c650922e120a3112a71be752dc8808d598bf Mon Sep 17 00:00:00 2001 From: Corey Goodfred Date: Fri, 3 May 2024 21:23:15 -0400 Subject: [PATCH] [COST-4744] Azure Data Transfer Generator (#498) * [COST-4744] Add Azure Data Transfer Generator to nise --- nise/__init__.py | 2 +- nise/generators/azure/__init__.py | 1 + nise/generators/azure/azure_generator.py | 5 +- .../azure/data_transfer_generator.py | 81 +++++++++++++++++++ nise/report.py | 2 + tests/test_azure_generator.py | 73 +++++++++++++++++ 6 files changed, 161 insertions(+), 3 deletions(-) create mode 100644 nise/generators/azure/data_transfer_generator.py diff --git a/nise/__init__.py b/nise/__init__.py index 13044923..7c79fbeb 100644 --- a/nise/__init__.py +++ b/nise/__init__.py @@ -1,3 +1,3 @@ -__version__ = "4.4.16" +__version__ = "4.4.17" VERSION = __version__.split(".") diff --git a/nise/generators/azure/__init__.py b/nise/generators/azure/__init__.py index 15b5c31c..8ecfce26 100644 --- a/nise/generators/azure/__init__.py +++ b/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 diff --git a/nise/generators/azure/azure_generator.py b/nise/generators/azure/azure_generator.py index 8044566a..327ba5eb 100644 --- a/nise/generators/azure/azure_generator.py +++ b/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 diff --git a/nise/generators/azure/data_transfer_generator.py b/nise/generators/azure/data_transfer_generator.py new file mode 100644 index 00000000..a6f51e2a --- /dev/null +++ b/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 . +# +"""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}") diff --git a/nise/report.py b/nise/report.py index e08f44a7..fa6ddfe2 100644 --- a/nise/report.py +++ b/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 diff --git a/tests/test_azure_generator.py b/tests/test_azure_generator.py index 8c0025d0..e7e0a261 100644 --- a/tests/test_azure_generator.py +++ b/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")