[azure][feat]: Improved AzureComputeDiskTypePricing deletion (#2231)

someengineering · Oct 8, 2024 · 8482aa6 · 8482aa6
1 parent 5f83fed
commit 8482aa6
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 21 deletions.
diff --git a/plugins/azure/fix_plugin_azure/collector.py b/plugins/azure/fix_plugin_azure/collector.py
@@ -276,20 +276,10 @@ def remove_usage_zero_value() -> None:
         rm_leaf_nodes(AzurePostgresqlServerType, AzureSubscription)
         rm_leaf_nodes(AzureCosmosDBLocation, AzureLocation, check_pred=False)
         rm_leaf_nodes(AzureLocation, check_pred=False)
+        rm_leaf_nodes(AzureComputeDiskTypePricing, AzureSubscription)
         remove_usage_zero_value()
         self.graph.remove_recursively(builder.nodes(AzureLocation, lambda r: r.compute_region_in_use(builder) is False))
 
-    def after_collect(self, builder: GraphBuilder) -> None:
-        # Filter unnecessary nodes such as AzureComputeDiskTypePricing
-        nodes_to_remove = []
-        node_types = (AzureComputeDiskTypePricing,)
-
-        for node in self.graph.nodes:
-            if not isinstance(node, node_types):
-                continue
-            nodes_to_remove.append(node)
-        self._delete_nodes(nodes_to_remove)
-
     def _delete_nodes(self, nodes_to_delete: Any) -> None:
         removed = set()
         for node in nodes_to_delete:

diff --git a/plugins/azure/test/collector_test.py b/plugins/azure/test/collector_test.py
@@ -11,7 +11,6 @@
 from fix_plugin_azure.config import AzureCredentials, AzureConfig
 from fix_plugin_azure.resource.base import MicrosoftResource, AzureSubscription, GraphBuilder
 from fix_plugin_azure.resource.compute import (
-    AzureComputeDiskTypePricing,
     AzureComputeVirtualMachine,
     AzureComputeVirtualMachineSize,
     AzureComputeDisk,
@@ -65,8 +64,6 @@ def test_filter(credentials: AzureCredentials, builder: GraphBuilder) -> None:
         AzureComputeVirtualMachineSize.collect(raw=json.load(f)["value"], builder=builder)
     with open(os.path.dirname(__file__) + "/files/compute/virtualMachines.json") as f:
         AzureComputeVirtualMachine.collect(raw=json.load(f)["value"], builder=builder)
-    with open(os.path.dirname(__file__) + "/files/compute/calculator.json") as f:
-        AzureComputeDiskTypePricing.collect(raw=json.load(f), builder=builder)
 
     collector = collector_with_graph(builder.graph, credentials)
 
@@ -78,13 +75,6 @@ def test_filter(credentials: AzureCredentials, builder: GraphBuilder) -> None:
         num_all_virtual_machine_types
     )
 
-    pricing_info = list(collector.graph.search("kind", "azure_compute_disk_type_pricing"))
-
-    assert len(pricing_info) > 0
-
-    collector.after_collect(builder)
-    assert len(list(collector.graph.search("kind", "azure_compute_disk_type_pricing"))) < len(pricing_info)
-
 
 def test_collect_cost(credentials: AzureCredentials, builder: GraphBuilder) -> None:
     with open(os.path.dirname(__file__) + "/files/compute/vmSizes.json") as f: