fix: Properly manage include in pre_tasks and post_tasks (#97)

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

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)
 

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

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

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

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)