fix: Updated tests and connecting resources stage

plugins/azure/resoto_plugin_azure/collector.py

@@ -115,9 +115,7 @@ def rm_nodes(cls, ignore_kinds: Optional[Type[Any]] = None) -> None:  # type: ig
                 pred = list(self.graph.predecessors(node))
                 if ignore_kinds is not None:
                     pred = [p for p in pred if not isinstance(p, ignore_kinds)]
-
                 if not pred:
-                    remove_nodes.extend(pred)
                     remove_nodes.append(node)
             removed = set()
             for node in remove_nodes:

plugins/azure/resoto_plugin_azure/resource/base.py

@@ -4,7 +4,7 @@
 from concurrent.futures import Future
 from datetime import datetime
 from threading import Lock
-from typing import Any, ClassVar, Dict, Optional, TypeVar, List, Tuple, Type, Callable, Union, cast
+from typing import Any, ClassVar, Dict, Optional, TypeVar, List, Type, Callable, cast
 
 from attr import define, field
 from azure.core.utils import CaseInsensitiveDict
@@ -134,40 +134,6 @@ def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         # Default behavior: add resource to the namespace
         pass
 
-    def fetch_resources(
-        self,
-        builder: GraphBuilder,
-        service: str,
-        api_version: str,
-        path: str,
-        path_parameters: List[str],
-        query_parameters: List[str],
-        first_property: Callable[[Json], Union[List[str], str]],
-        second_property: Callable[[Json], Union[List[str], str]],
-    ) -> List[Tuple[Union[str, List[str]], Union[str, List[str]]]]:
-        """
-        Fetch additional resources from the Azure API.
-        Can used for further connection using the connect_in_graph method.
-
-        Parameters:
-        - first_property: Compared property.
-        - second_property: Binding property | Compared property.
-
-        Returns:
-        List[Tuple[Union[str, List[str]], Union[str, List[str]]]]: A list of tuples containing information to compare and connect the retrieved resources.
-        """
-        resources_api_spec = AzureApiSpec(
-            service=service,
-            version=api_version,
-            path=path,
-            path_parameters=path_parameters,
-            query_parameters=query_parameters,
-            access_path="value",
-            expect_array=True,
-        )
-
-        return [(first_property(r), second_property(r)) for r in builder.client.list(resources_api_spec)]
-
     @classmethod
     def collect_resources(
         cls: Type[AzureResourceType], builder: GraphBuilder, **kwargs: Any

plugins/azure/resoto_plugin_azure/resource/network.py

@@ -2822,21 +2822,18 @@ class AzureDscpConfiguration(AzureResource):
     source_port_ranges: Optional[List[AzureQosPortRange]] = field(default=None, metadata={'description': 'Sources port ranges.'})  # fmt: skip
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
-        second_property: Callable[[Json], List[str]] = lambda r: [  # pylint: disable=unnecessary-lambda-assignment
-            ip_config["properties"]["subnet"]["id"]
-            for ip_config in r.get("properties", {}).get("ipConfigurations", [])
-            if "subnet" in ip_config.get("properties", {})
+        first_property: Callable[[AzureNetworkInterface], str] = lambda n: n.id or ""
+
+        second_property: Callable[[AzureNetworkInterface], List[str]] = lambda n: [
+            ip_config._subnet_id
+            for ip_config in n.interface_ip_configurations or []
+            if ip_config._subnet_id is not None
         ]
-        nis_and_subnet_id = self.fetch_resources(
-            builder,
-            service="network",
-            api_version="2023-05-01",
-            path="/subscriptions/{subscriptionId}/providers/Microsoft.Network/networkInterfaces",
-            path_parameters=["subscriptionId"],
-            query_parameters=["api-version"],
-            first_property=lambda r: r["id"],
-            second_property=second_property,
-        )
+
+        nis_and_subnet_id = [
+            (first_property(n), second_property(n)) for n in builder.nodes(clazz=AzureNetworkInterface)
+        ]
+
         if (network_interfaces := self._associated_network_interface_ids) and (ni_ids_and_s_id := nis_and_subnet_id):
             for network_interface_id in network_interfaces:
                 for info in ni_ids_and_s_id:
@@ -3116,16 +3113,13 @@ class AzureExpressRouteCircuit(AzureResource):
     stag: Optional[int] = field(default=None, metadata={'description': 'The identifier of the circuit traffic. Outer tag for QinQ encapsulation.'})  # fmt: skip
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
-        ids_and_names_in_resource = self.fetch_resources(
-            builder,
-            service="network",
-            api_version="2023-05-01",
-            path="/subscriptions/{subscriptionId}/providers/Microsoft.Network/ExpressRoutePortsLocations",
-            path_parameters=["subscriptionId"],
-            query_parameters=["api-version"],
-            first_property=lambda r: r["name"],
-            second_property=lambda r: r["id"],
-        )
+        first_property: Callable[[AzureExpressRoutePortsLocation], str] = lambda n: n.name or ""
+
+        second_property: Callable[[AzureExpressRoutePortsLocation], str] = lambda n: n.id or ""
+
+        ids_and_names_in_resource = [
+            (first_property(n), second_property(n)) for n in builder.nodes(clazz=AzureExpressRoutePortsLocation)
+        ]
 
         if route_port_id := self.express_route_port:
             builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureExpressRoutePort, id=route_port_id)
@@ -3802,17 +3796,14 @@ class AzureIpGroup(AzureResource):
     provisioning_state: Optional[str] = field(default=None, metadata={'description': 'The current provisioning state.'})  # fmt: skip
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
-        virtual_networks = self.fetch_resources(
-            builder,
-            service="network",
-            api_version="2023-05-01",
-            path="/subscriptions/{subscriptionId}/providers/Microsoft.Network/virtualNetworks",
-            path_parameters=["subscriptionId"],
-            query_parameters=["api-version"],
-            first_property=lambda r: r["properties"]["addressSpace"]["addressPrefixes"],
-            second_property=lambda r: r["id"],
+        first_property: Callable[[AzureVirtualNetwork], List[str]] = (
+            lambda n: n.address_space.address_prefixes if n.address_space and n.address_space.address_prefixes else []
         )
 
+        second_property: Callable[[AzureVirtualNetwork], str] = lambda n: n.id or ""
+
+        virtual_networks = [(first_property(n), second_property(n)) for n in builder.nodes(clazz=AzureVirtualNetwork)]
+
         if (ip_addresses := self.ip_addresses) and (vns := virtual_networks):
             for ip_address in ip_addresses:
                 for info in vns:
@@ -4220,22 +4211,15 @@ def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         # Import placed inside the method due to circular import error resolution
         from resoto_plugin_azure.resource.compute import AzureVirtualMachine  # pylint: disable=import-outside-toplevel
 
-        first_property: Callable[[Json], List[str]] = lambda r: [  # pylint: disable=unnecessary-lambda-assignment
-            ip_config["id"] for ip_config in r.get("properties", {}).get("ipConfigurations", [])
+        first_property: Callable[[AzureNetworkInterface], List[str]] = lambda n: [
+            ip_config.id for ip_config in n.interface_ip_configurations or [] if ip_config.id is not None
         ]
 
-        ip_confs_and_vm_ids = self.fetch_resources(
-            builder,
-            service="network",
-            api_version="2023-05-01",
-            path="/subscriptions/{subscriptionId}/providers/Microsoft.Network/networkInterfaces",
-            path_parameters=["subscriptionId"],
-            query_parameters=["api-version"],
-            first_property=first_property,
-            second_property=lambda r: r["properties"]["virtualMachine"]["id"]
-            if "properties" in r and "virtualMachine" in r["properties"]
-            else "",
-        )
+        second_property: Callable[[AzureNetworkInterface], str] = lambda n: n.virtual_machine or ""
+
+        ip_confs_and_vm_ids = [
+            (first_property(n), second_property(n)) for n in builder.nodes(clazz=AzureNetworkInterface)
+        ]
 
         if container_nic := self.container_network_interface_configurations:
             for container in container_nic:
@@ -4383,16 +4367,13 @@ class AzureNetworkVirtualAppliance(AzureResource):
     virtual_hub: Optional[str] = field(default=None, metadata={"description": "Reference to another subresource."})
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
-        vendors_in_resource = self.fetch_resources(
-            builder,
-            service="network",
-            api_version="2023-05-01",
-            path="/subscriptions/{subscriptionId}/providers/Microsoft.Network/networkVirtualApplianceSkus",
-            path_parameters=["subscriptionId"],
-            query_parameters=["api-version"],
-            first_property=lambda r: r["properties"]["vendor"],
-            second_property=lambda r: r["name"],
-        )
+        first_property: Callable[[AzureNetworkVirtualApplianceSku], str] = lambda n: n.vendor or ""
+
+        second_property: Callable[[AzureNetworkVirtualApplianceSku], str] = lambda n: n.name or ""
+
+        vendors_in_resource = [
+            (first_property(n), second_property(n)) for n in builder.nodes(clazz=AzureNetworkVirtualApplianceSku)
+        ]
 
         if (nva := self.nva_sku) and (nva_vendor := nva.vendor) and (vendors := vendors_in_resource):
             for info in vendors:
@@ -4473,16 +4454,13 @@ class AzureNetworkWatcher(AzureResource):
     location: Optional[str] = field(default=None, metadata={"description": "Resource location."})
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
-        locations_and_ids_in_vn = self.fetch_resources(
-            builder,
-            service="network",
-            api_version="2023-05-01",
-            path="/subscriptions/{subscriptionId}/providers/Microsoft.Network/virtualNetworks",
-            path_parameters=["subscriptionId"],
-            query_parameters=["api-version"],
-            first_property=lambda r: r["location"],
-            second_property=lambda r: r["id"],
-        )
+        first_property: Callable[[AzureVirtualNetwork], str] = lambda n: n.location or ""
+
+        second_property: Callable[[AzureVirtualNetwork], str] = lambda n: n.id or ""
+
+        locations_and_ids_in_vn = [
+            (first_property(n), second_property(n)) for n in builder.nodes(clazz=AzureVirtualNetwork)
+        ]
 
         if (nw_location := self.location) and (vns_info := locations_and_ids_in_vn):
             for info in vns_info:
@@ -4950,25 +4928,19 @@ class AzureVirtualHub(AzureResource):
     vpn_gateway: Optional[str] = field(default=None, metadata={"description": "Reference to another subresource."})
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
-        first_property: Callable[[Json], List[str]] = lambda r: [
-            ip_config["id"] for ip_config in r.get("properties", {}).get("ipConfigurations", [])
+        first_property: Callable[[AzureNetworkInterface], List[str]] = lambda n: [
+            ip_config.id for ip_config in n.interface_ip_configurations or [] if ip_config.id is not None
         ]
-        second_property: Callable[[Json], List[str]] = lambda r: [
-            ip_config["properties"]["publicIPAddress"]["id"]
-            for ip_config in r.get("properties", {}).get("ipConfigurations", [])
-            if "publicIPAddress" in ip_config.get("properties", {})
+
+        second_property: Callable[[AzureNetworkInterface], List[str]] = lambda n: [
+            ip_config._public_ip_id
+            for ip_config in n.interface_ip_configurations or []
+            if ip_config._public_ip_id is not None
         ]
 
-        p_ip_addresses_ip_c_ids = self.fetch_resources(
-            builder,
-            service="network",
-            api_version="2023-05-01",
-            path="/subscriptions/{subscriptionId}/providers/Microsoft.Network/networkInterfaces",
-            path_parameters=["subscriptionId"],
-            query_parameters=["api-version"],
-            first_property=first_property,
-            second_property=second_property,
-        )
+        p_ip_a_and_ip_c_ids = [
+            (first_property(n), second_property(n)) for n in builder.nodes(clazz=AzureNetworkInterface)
+        ]
 
         if er_gateway_id := self.express_route_gateway:
             builder.add_edge(
@@ -4978,7 +4950,7 @@ def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
             builder.add_edge(self, edge_type=EdgeType.default, reverse=True, clazz=AzureVpnGateway, id=vpn_gateway_id)
         if vw_id := self.virtual_wan:
             builder.add_edge(self, edge_type=EdgeType.default, reverse=True, clazz=AzureVirtualWAN, id=vw_id)
-        if (ip_config_ids := self.ip_configuration_ids) and (p_ip_a_and_ip_conf_ids := p_ip_addresses_ip_c_ids):
+        if (ip_config_ids := self.ip_configuration_ids) and (p_ip_a_and_ip_conf_ids := p_ip_a_and_ip_c_ids):
             for ip_config_id in ip_config_ids:
                 for info in p_ip_a_and_ip_conf_ids:
                     collected_ip_conf_ids, p_ip_address_ids = info
@@ -5101,6 +5073,7 @@ class AzureVirtualNetwork(AzureResource, BaseNetwork):
         "_subnet_ids": S("properties", "subnets", default=[]) >> ForallBend(S("id")),
         "virtual_network_peerings": S("properties", "virtualNetworkPeerings")
         >> ForallBend(AzureVirtualNetworkPeering.mapping),
+        "location": S("location"),
     }
     address_space: Optional[AzureAddressSpace] = field(default=None, metadata={'description': 'AddressSpace contains an array of IP address ranges that can be used by subnets of the virtual network.'})  # fmt: skip
     bgp_communities: Optional[AzureVirtualNetworkBgpCommunities] = field(default=None, metadata={'description': 'Bgp Communities sent over ExpressRoute with each route corresponding to a prefix in this VNET.'})  # fmt: skip
@@ -5118,6 +5091,7 @@ class AzureVirtualNetwork(AzureResource, BaseNetwork):
     resource_guid: Optional[str] = field(default=None, metadata={'description': 'The resourceGuid property of the Virtual Network resource.'})  # fmt: skip
     _subnet_ids: Optional[List[str]] = field(default=None, metadata={'description': 'A list of subnets in a Virtual Network.'})  # fmt: skip
     virtual_network_peerings: Optional[List[AzureVirtualNetworkPeering]] = field(default=None, metadata={'description': 'A list of peerings in a Virtual Network.'})  # fmt: skip
+    location: Optional[str] = field(default=None, metadata={"description": "Resource location."})
 
     def post_process(self, graph_builder: GraphBuilder, source: Json) -> None:
         def collect_subnets() -> None:

plugins/azure/test/collector_test.py

@@ -1,9 +1,30 @@
+import os
+import json
+from queue import Queue
+
 from resoto_plugin_azure.azure_client import AzureClient
 from resoto_plugin_azure.collector import AzureSubscriptionCollector
 from resoto_plugin_azure.config import AzureCredentials, AzureConfig
-from resoto_plugin_azure.resource.base import AzureSubscription
+from resoto_plugin_azure.resource.base import AzureSubscription, GraphBuilder
+from resoto_plugin_azure.resource.compute import AzureVirtualMachine, AzureVirtualMachineSize
 from resotolib.baseresources import Cloud
 from resotolib.core.actions import CoreFeedback
+from resotolib.graph import Graph
+
+
+def collector_with_graph(
+    graph: Graph,
+    credentials: AzureCredentials,
+) -> AzureSubscriptionCollector:
+    collector = AzureSubscriptionCollector(
+        config=AzureConfig(),
+        cloud=Cloud(id="azure"),
+        subscription=AzureSubscription(id="test", subscription_id="test"),
+        credentials=credentials,
+        core_feedback=CoreFeedback("test", "test", "test", Queue()),
+    )
+    collector.graph = graph
+    return collector
 
 
 def test_collect(
@@ -17,3 +38,18 @@ def test_collect(
     collector.collect()
     assert len(collector.graph.nodes) == 418
     assert len(collector.graph.edges) == 473
+
+
+def test_filter(credentials: AzureCredentials, builder: GraphBuilder) -> None:
+    with open(os.path.dirname(__file__) + "/files/compute/vmSizes.json") as f:
+        AzureVirtualMachineSize.collect(raw=json.load(f)["value"], builder=builder)
+    with open(os.path.dirname(__file__) + "/files/compute/virtualMachines.json") as f:
+        AzureVirtualMachine.collect(raw=json.load(f)["value"], builder=builder)
+
+    collector = collector_with_graph(builder.graph, credentials)
+
+    num_all_virtual_machine_types = list(collector.graph.search("kind", "azure_virtual_machine_size"))
+
+    collector.filter_nodes()
+
+    assert len(list(collector.graph.search("kind", "azure_virtual_machine_size"))) < len(num_all_virtual_machine_types)

plugins/azure/test/compute_test.py

@@ -134,6 +134,7 @@ def test_virtual_machine(builder: GraphBuilder) -> None:
     assert len(builder.edges_of(AzureVirtualMachine, AzureNetworkInterface)) == 1
     assert len(builder.edges_of(AzureNetworkSecurityGroup, AzureVirtualMachine)) == 1
     assert len(builder.edges_of(AzureLoadBalancer, AzureVirtualMachine)) == 1
+    assert len(builder.edges_of(AzureVirtualMachine, AzureVirtualMachineSize)) == 2
 
 
 def test_virtual_machine_resources(builder: GraphBuilder) -> None:

plugins/azure/test/conftest.py

@@ -96,8 +96,10 @@ def builder(
         executor=executor_queue,
         core_feedback=core_feedback,
     )
-    location = AzureLocation(id="westeurope", display_name="West Europe")
-    builder.location_lookup = {"westeurope": location}
+    location_west = AzureLocation(id="westeurope", display_name="West Europe", name="westeurope")
+    location_east = AzureLocation(id="eastus", display_name="East US", name="eastus")
+    builder.location_lookup = {"westeurope": location_west, "eastus": location_east}
+    builder.location = location_east
     return builder