def prepare_graph(do_client) -> Graph: cloud = Cloud("do") team = DigitalOceanTeam(id="test_team", urn="do:team:test_team") plugin_instance = DigitalOceanTeamCollector(team, do_client) plugin_instance.collect() cloud_graph = Graph(root=cloud) graph = Graph(root=GraphRoot("root", {})) cloud_graph.merge(plugin_instance.graph) graph.merge(cloud_graph) sanitize(graph) return graph
def collect(collectors: List[BaseCollectorPlugin]) -> Graph: graph = Graph(root=GraphRoot("root", {})) max_workers = (len(collectors) if len(collectors) < self._config.resotoworker.pool_size else self._config.resotoworker.pool_size) if max_workers == 0: log.error( "No workers configured or no collector plugins loaded - skipping collect" ) return pool_args = {"max_workers": max_workers} if self._config.resotoworker.fork_process: pool_args["mp_context"] = multiprocessing.get_context("spawn") pool_args["initializer"] = resotolib.proc.initializer pool_executor = futures.ProcessPoolExecutor collect_args = { "args": ArgumentParser.args, "running_config": self._config.running_config, } else: pool_executor = futures.ThreadPoolExecutor collect_args = {} with pool_executor(**pool_args) as executor: wait_for = [ executor.submit( collect_plugin_graph, collector, **collect_args, ) for collector in collectors ] for future in futures.as_completed(wait_for): cluster_graph = future.result() if not isinstance(cluster_graph, Graph): log.error( f"Skipping invalid cluster_graph {type(cluster_graph)}" ) continue graph.merge(cluster_graph) sanitize(graph) return graph
def test_graph_merge(): rg1 = Graph() rg2 = Graph() a = SomeTestResource("a", {}) b = SomeTestResource("b", {}) c = SomeTestResource("c", {}) d = SomeTestResource("d", {}) rg1.add_node(a) rg1.add_node(b) rg2.add_node(c) rg2.add_node(d) rg1.add_edge(a, b, edge_type=EdgeType.delete) rg2.add_edge(c, d, edge_type=EdgeType.delete) rg1.merge(rg2) assert len(rg1.nodes) == 4 assert len(rg1.edges) == 2 for edge in rg1.edges: assert len(edge) == 3 key = edge[2] assert len(key) == 3 edge_type = key[2] assert edge_type == EdgeType.delete