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
