From 1f2f2be40860647ecb5047db9fc3b86b951c054c Mon Sep 17 00:00:00 2001 From: Mohamed El Mouctar Haidara Date: Thu, 6 Jan 2022 23:49:52 +0100 Subject: [PATCH] fix: Properly manage include in pre_tasks and post_tasks (#97) --- ansibleplaybookgrapher/graph.py | 80 +++++++++++++++++++++++------- ansibleplaybookgrapher/renderer.py | 2 +- tests/fixtures/with_block.yml | 3 ++ tests/test_graph.py | 36 +++++++++++++- tests/test_parser.py | 34 +++++++++++-- tests/test_playbook_grapher.py | 7 +-- 6 files changed, 136 insertions(+), 26 deletions(-) diff --git a/ansibleplaybookgrapher/graph.py b/ansibleplaybookgrapher/graph.py index 28c610ca..59510366 100644 --- a/ansibleplaybookgrapher/graph.py +++ b/ansibleplaybookgrapher/graph.py @@ -1,5 +1,5 @@ from collections import defaultdict -from typing import Dict, List +from typing import Dict, List, ItemsView from ansibleplaybookgrapher.utils import generate_id @@ -41,12 +41,23 @@ def __init__(self, node_name: str, node_id: str, raw_object=None, supported_comp :param node_name: :param node_id: + :param raw_object: The raw ansible object matching this node in the graph. Will be None if there is no match on + Ansible side + :param supported_compositions: """ super().__init__(node_name, node_id, raw_object) self._supported_compositions = supported_compositions or [] # The dict will contain the different types of composition. self._compositions = defaultdict(list) # type: Dict[str, List] + def items(self) -> ItemsView[str, List[Node]]: + """ + Return a view object (list of tuples) of all the nodes inside this composite node. The first element of the + tuple is the composition name and the second one a list of nodes + :return: + """ + return self._compositions.items() + def add_node(self, target_composition: str, node: Node): """ Add a node in the target composition @@ -61,8 +72,8 @@ def add_node(self, target_composition: str, node: Node): def links_structure(self) -> Dict[Node, List[Node]]: """ - Return a representation of the composite node where each key of the dictionary is the node ID and the values is - a list of the linked nodes + Return a representation of the composite node where each key of the dictionary is the node and the value is the + list of the linked nodes :return: """ links = defaultdict(list) @@ -81,6 +92,25 @@ def _get_all_links(self, links: Dict[Node, List[Node]]): links[self].append(node) +class CompositeTasksNode(CompositeNode): + """ + A special composite node which only support adding "tasks" + """ + + def __init__(self, node_name: str, node_id: str, raw_object=None): + super().__init__(node_name, node_id, raw_object=raw_object) + self._supported_compositions = ["tasks"] + + def add_node(self, target_composition: str, node: Node): + """ + Override the add_node because block only contains "tasks" regardless of the context (pre_tasks or post_tasks) + :param target_composition: This is ignored. It's always "tasks" for block + :param node: + :return: + """ + super().add_node("tasks", node) + + class PlaybookNode(CompositeNode): """ A playbook is a list of play @@ -146,14 +176,13 @@ def tasks(self) -> List['EdgeNode']: return self._compositions["tasks"] -class BlockNode(CompositeNode): +class BlockNode(CompositeTasksNode): """ A block node: https://docs.ansible.com/ansible/latest/user_guide/playbooks_blocks.html """ def __init__(self, node_name: str, node_id: str = None, raw_object=None): - super().__init__(node_name, node_id or generate_id("block_"), raw_object=raw_object, - supported_compositions=["tasks"]) + super().__init__(node_name, node_id or generate_id("block_"), raw_object=raw_object) @property def tasks(self) -> List['EdgeNode']: @@ -163,15 +192,6 @@ def tasks(self) -> List['EdgeNode']: """ return self._compositions['tasks'] - def add_node(self, target_composition: str, node: Node): - """ - Override the add_node because block only contains "tasks" regardless of the context (pre_tasks or post_tasks) - :param target_composition: This is ignored. It's always "tasks" for block - :param node: - :return: - """ - super().add_node("tasks", node) - class EdgeNode(CompositeNode): """ @@ -221,14 +241,40 @@ def __init__(self, node_name: str, node_id: str = None, raw_object=None): super().__init__(node_name, node_id or generate_id("task_"), raw_object) -class RoleNode(CompositeNode): +class RoleNode(CompositeTasksNode): """ A role node. A role is a composition of tasks """ def __init__(self, node_name: str, node_id: str = None, raw_object=None): - super().__init__(node_name, node_id or generate_id("role_"), raw_object, supported_compositions=["tasks"]) + super().__init__(node_name, node_id or generate_id("role_"), raw_object=raw_object) @property def tasks(self): return self._compositions["tasks"] + + +def _get_all_tasks_nodes(composite: CompositeNode, task_acc: List[TaskNode]): + """ + :param composite: + :param task_acc: + :return: + """ + items = composite.items() + for _, nodes in items: + for node in nodes: + if isinstance(node, TaskNode): + task_acc.append(node) + elif isinstance(node, CompositeNode): + _get_all_tasks_nodes(node, task_acc) + + +def get_all_tasks_nodes(composite: CompositeNode) -> List[TaskNode]: + """ + Return all the TaskNode inside a composite node + :param composite: + :return: + """ + tasks = [] + _get_all_tasks_nodes(composite, tasks) + return tasks diff --git a/ansibleplaybookgrapher/renderer.py b/ansibleplaybookgrapher/renderer.py index 4d43c3fc..8102cd5e 100644 --- a/ansibleplaybookgrapher/renderer.py +++ b/ansibleplaybookgrapher/renderer.py @@ -53,7 +53,7 @@ def render_node(self, graph: Digraph, edge: EdgeNode, color: str, node_counter: else: edge_label = f"{node_counter} {edge.name}" graph.node(destination_node.id, label=node_label_prefix + destination_node.name, shape=shape, - id=destination_node.id, tooltip=destination_node.name, color=color) + id=destination_node.id, tooltip=destination_node.name, color=color, hello=destination_node.name) graph.edge(source_node.id, destination_node.id, label=edge_label, color=color, fontcolor=color, id=edge.id, tooltip=edge_label, labeltooltip=edge_label) diff --git a/tests/fixtures/with_block.yml b/tests/fixtures/with_block.yml index 0a2f90e0..5caba63b 100644 --- a/tests/fixtures/with_block.yml +++ b/tests/fixtures/with_block.yml @@ -1,6 +1,9 @@ --- - hosts: all pre_tasks: + - name: Include role + include_role: + name: fake_role - name: Block in pre task block: - name: debug diff --git a/tests/test_graph.py b/tests/test_graph.py index 0d1a6a49..25596933 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -1,4 +1,4 @@ -from ansibleplaybookgrapher.graph import RoleNode, TaskNode, EdgeNode, PlayNode +from ansibleplaybookgrapher.graph import RoleNode, TaskNode, EdgeNode, PlayNode, BlockNode, get_all_tasks_nodes def test_links_structure(): @@ -39,3 +39,37 @@ def test_links_structure(): for e in [edge_1, edge_2, edge_3, edge_role]: assert len(all_links[e]) == 1, "An edge should be linked to one node" + + +def test_get_all_tasks_nodes(): + """ + Test the function get_all_tasks_nodes + :return: + """ + play = PlayNode("play") + + role_1 = RoleNode("my_role_1") + edge_role = EdgeNode(play, role_1, "from play to role") + play.add_node("roles", edge_role) + + # play -> role 1 -> edge 1 -> task 1 + task_1 = TaskNode("task 1") + edge_1 = EdgeNode(role_1, task_1, "from role1 to task 1") + role_1.add_node("tasks", edge_1) + + # play -> block_1 -> task 2 and task 3 + block_1 = BlockNode("block 1") + task_2 = TaskNode("task 2") + task_3 = TaskNode("task 3") + block_1.add_node("tasks", task_2) + block_1.add_node("tasks", task_3) + play.add_node("tasks", block_1) + # play -> block_1 -> block_2 -> task 4 + block_2 = BlockNode("block 2") + task_4 = TaskNode("task 4") + block_2.add_node("tasks", task_4) + block_1.add_node("tasks", block_2) + + all_tasks = get_all_tasks_nodes(play) + assert len(all_tasks) == 4, "There should be 4 tasks in all" + assert [task_1, task_2, task_3, task_4] == all_tasks diff --git a/tests/test_parser.py b/tests/test_parser.py index 5a607272..651c2665 100644 --- a/tests/test_parser.py +++ b/tests/test_parser.py @@ -1,9 +1,25 @@ +from typing import List + import pytest from ansible.utils.display import Display from ansibleplaybookgrapher import PlaybookParser from ansibleplaybookgrapher.cli import PlaybookGrapherCLI -from ansibleplaybookgrapher.graph import TaskNode, BlockNode, RoleNode +from ansibleplaybookgrapher.graph import TaskNode, BlockNode, RoleNode, get_all_tasks_nodes, CompositeNode + + +def get_all_tasks(composites: List[CompositeNode]) -> List[TaskNode]: + """ + + :param composites: + :return: + """ + tasks = [] + + for c in composites: + tasks.extend(get_all_tasks_nodes(c)) + + return tasks @pytest.mark.parametrize('grapher_cli', [["include_role.yml"]], indirect=True) @@ -48,15 +64,25 @@ def test_block_parsing(grapher_cli: PlaybookGrapherCLI, display: Display): :param display: :return: """ - parser = PlaybookParser(grapher_cli.options.playbook_filename, display=display) + parser = PlaybookParser(grapher_cli.options.playbook_filename, display=display, include_role_tasks=True) playbook_node = parser.parse() assert len(playbook_node.plays) == 1 play_node = playbook_node.plays[0].destination + pre_tasks = play_node.pre_tasks tasks = play_node.tasks post_tasks = play_node.post_tasks - assert len(tasks) == 3 - assert len(post_tasks) == 2 + total_pre_tasks = get_all_tasks(pre_tasks) + total_tasks = get_all_tasks(tasks) + total_post_tasks = get_all_tasks(post_tasks) + assert len(total_pre_tasks) == 4, f"The play should contain 4 pre tasks but we found {len(total_pre_tasks)} pre task(s)" + assert len(total_tasks) == 7, f"The play should contain 3 tasks but we found {len(total_tasks)} task(s)" + assert len( + total_post_tasks) == 2, f"The play should contain 2 post tasks but we found {len(total_post_tasks)} post task(s)" + + # Check pre tasks + assert isinstance(pre_tasks[0].destination, RoleNode), "The first edge should have a RoleNode as destination" + assert isinstance(pre_tasks[1].destination, BlockNode), "The second edge should have a BlockNode as destination" # Check tasks task_1 = tasks[0].destination diff --git a/tests/test_playbook_grapher.py b/tests/test_playbook_grapher.py index 1b13659d..75f74505 100644 --- a/tests/test_playbook_grapher.py +++ b/tests/test_playbook_grapher.py @@ -165,15 +165,16 @@ def test_with_block(request): """ Test with_roles.yml, an example with roles """ - svg_path, playbook_path = run_grapher("with_block.yml", output_filename=request.node.name) + svg_path, playbook_path = run_grapher("with_block.yml", output_filename=request.node.name, + additional_args=["--include-role-tasks"]) _common_tests(svg_path=svg_path, playbook_path=playbook_path, plays_number=1, tasks_number=7, post_tasks_number=2, - pre_tasks_number=1, blocks_number=4) + pre_tasks_number=4, blocks_number=4, roles_number=1) def test_nested_include_tasks(request): """ - Test nested_include.yml, an example with an include tasks that include another tasks + Test nested_include.yml, an example with an include_tasks that include another tasks """ svg_path, playbook_path = run_grapher("nested_include_tasks.yml", output_filename=request.node.name)