Add tests

Author: FreeYourSoul

CHANGELOG.md

@@ -32,21 +32,21 @@
 
 ## Feature todo
 
-* Implement sub_graph
+* ~~Implement sub_graph~~
 * ~~fix graph node enforcement of uid~~
 * ~~Delete node~~ 
 * ~~Partial graph traversal~~
 
 ## Documentation todo
 
+* ~~Add documentation on sub_graph~~
 * ~~Add doc about node / graph creation~~
 * ~~Add documentation about graph nodes~~
 * ~~Add documentation about custom hook~~
 * ~~Add documentation Composed Visitor~~
-* Add documentation on sub_graph
 
 ## Testing todo
 
-* Test sub_graph
+* ~~Test sub_graph~~
 * ~~Test Composed Visitor~~
 * ~~Custom hook~~

README.md

@@ -19,6 +19,7 @@ It provide:
   * [Graph creation](#graph-creation)
   * [Graph Node](#graph-node) 
   * [Fork mechanism](#fork) 
+  * [Sub graph](#sub-graph) 
 * [Visitors](#visitors)
     * [Visitor hook](#abstract-visitor-hooks)
     * [Standard visitors](#standard-visitor-provided-by-freexgraph)
@@ -258,6 +259,20 @@ It is also possible to provide a join node. It will be a node used as join for t
 But if you want to do a map reduce with a fork, it is do-able by setting the join_id to the `fork_from_node` method. The join_id has to be an existing node on which, for every parents that are part of the fork has only this join node as child.
 See [test using this mechanism](https://github.com/FreeYourSoul/FreExGraph/blob/ae707cf0fcb8486bde783cd0c7fe67217a56b3d2/test/fork_test.py#L41-L66) for more details
 
+### Sub-Graph
+
+A method exist in the `FreExGraph` class in order to make sub graph out of the graph.
+The signature is the following:
+```python
+def sub_graph(
+    self, from_node_id: str, to_nodes_id: Optional[List[str]] = None
+) -> FreExGraph:
+```
+
+Providing a node to start the subgraph and optionally a set of node where you want the sub graph to end. If no node matching the ``to_nodes_id`` is encountered, the subgraph go until the leaf nodes of the graph.
+
+This feature may be used as a fork mechanism (as seen above). It is easier to manipulate and monitor.
+
 ## Visitors
 
 ### Abstract Visitor hooks

freexgraph/freexgraph.py

@@ -296,26 +296,29 @@ def sub_graph(
         ), f"Error sub graph from node {from_node_id}, node has to be in the execution graph"
 
         nodes_in_subgraph: List[FreExNode] = []
-        nodes_in_subgraph_id: List[str] = []
+        nodes_in_subgraph_id: Set[str] = set()
 
         def add_node_in_subgraph(current_node: FreExNode):
-            current_node.parents = {
-                p for p in current_node.parents if p in nodes_in_subgraph_id
-            }
+            if current_node.id in nodes_in_subgraph_id:
+                return
             nodes_in_subgraph.append(current_node)
-            if to_nodes_id is not None and current_node in to_nodes_id:
+            nodes_in_subgraph_id.add(current_node.id)
+            if to_nodes_id is not None and current_node.id in to_nodes_id:
                 return
             all_suc = list(self._graph.successors(current_node.id))
             for successor in all_suc:
-                n = self.get_node(successor)
+                node_suc = self.get_node(successor)
                 assert (
-                    n is not None
-                ), f"Error sub graph to node {n.id}, node has to be in the execution graph"
-                add_node_in_subgraph(n)
-                nodes_in_subgraph_id.append(n.id)
+                    node_suc is not None
+                ), f"Error sub graph to node {node_suc.id}, node has to be in the execution graph"
+                add_node_in_subgraph(node_suc)
 
         add_node_in_subgraph(from_node)
 
+        # cleanup parents
+        for n in nodes_in_subgraph:
+            n.parents = {p for p in n.parents if p in nodes_in_subgraph_id}
+
         sub_graph = FreExGraph()
         sub_graph.add_nodes(nodes_in_subgraph)
         return sub_graph

test/conftest.py

@@ -125,7 +125,7 @@ def unordered_node_list_for_complex_graph() -> List[NodeForTest]:
     #
     return [
         NodeForTest("C", parents={"A"}),
-        NodeForTest("K", parents={"G"}),
+        NodeForTest("K", parents={"G", "F"}),
         NodeForTest("M", parents={"K", "B"}),
         NodeForTest("D", parents={"A"}),
         NodeForTest("J", parents={"F"}),

test/sub_graph_test.py

@@ -23,9 +23,9 @@
 
 import pytest
 
-from freexgraph import GraphNode, FreExNode, FreExGraph
+from freexgraph import FreExGraph
 
-# Graph used in thoses tests are from valid_complex_graph fixture
+# Graph used in theses tests are from valid_complex_graph fixture
 #
 #
 #            A                      B
@@ -51,30 +51,108 @@ def test_sub_graph_without_to_node(valid_complex_graph, visitor_test):
 
     sub_graph: FreExGraph = valid_complex_graph.sub_graph(from_node_id="F")
 
-    visitor_test.visit(valid_complex_graph)
     visitor_test.visit(sub_graph.root)
+
     assert len(visitor_test.visited) == 7
+    sorted_visit = sorted(visitor_test.visited)
+
+    assert sorted_visit[0] == "F"
+    assert sorted_visit[1] == "H"
+    assert sorted_visit[2] == "I"
+    assert sorted_visit[3] == "J"
+    assert sorted_visit[4] == "K"
+    assert sorted_visit[5] == "L"
+    assert sorted_visit[6] == "M"
+
+
+def test_sub_graph_of_one_element(valid_complex_graph, visitor_test):
+    sub_graph: FreExGraph = valid_complex_graph.sub_graph(
+        from_node_id="E", to_nodes_id=["E"]
+    )
+
+    visitor_test.visit(sub_graph.root)
 
+    assert len(visitor_test.visited) == 1
+    assert visitor_test.visited[0] == "E"
 
-def test_sub_graph_of_one_element(valid_complex_graph):
-    ...
 
+def test_sub_graph_with_to_node(valid_complex_graph, visitor_test):
+    # From A to K and J (L and M are not taken as K and J are stopping the subgraph)
+    #
+    #            A
+    #         /     \
+    #        C       D
+    #                 \
+    #                  F .______,
+    #               /  |  \     \
+    #             H    I   J     `K
+    #
+    sub_graph: FreExGraph = valid_complex_graph.sub_graph(
+        from_node_id="A", to_nodes_id=["K", "J"]
+    )
+
+    visitor_test.visit(sub_graph.root)
 
-def test_sub_graph_with_to_node_not_found(valid_complex_graph):
-    ...
+    assert len(visitor_test.visited) == 8
+    sorted_visit = sorted(visitor_test.visited)
 
+    assert sorted_visit[0] == "A"
+    assert sorted_visit[1] == "C"
+    assert sorted_visit[2] == "D"
+    assert sorted_visit[3] == "F"
+    assert sorted_visit[4] == "H"
+    assert sorted_visit[5] == "I"
+    assert sorted_visit[6] == "J"
+    assert sorted_visit[7] == "K"
 
-def test_sub_graph_with_to_node(valid_complex_graph):
-    ...
 
+def test_sub_graph_with_to_node_not_found(valid_complex_graph, visitor_test):
+    # From F to A, A is not in the sub graph starting from F, so it's like no to_node where set
+    #
+    #                  F .______,
+    #               /  |  \     \
+    #             H    I   J     `,K.
+    #                     /,_____/   \
+    #                    L             M
+
+    sub_graph: FreExGraph = valid_complex_graph.sub_graph(from_node_id="F", to_nodes_id="A")
+
+    visitor_test.visit(sub_graph.root)
+
+    assert len(visitor_test.visited) == 7
+    sorted_visit = sorted(visitor_test.visited)
 
-def test_sub_graph_with_to_node_partial(valid_complex_graph):
-    ...
+    assert sorted_visit[0] == "F"
+    assert sorted_visit[1] == "H"
+    assert sorted_visit[2] == "I"
+    assert sorted_visit[3] == "J"
+    assert sorted_visit[4] == "K"
+    assert sorted_visit[5] == "L"
+    assert sorted_visit[6] == "M"
 
 
 def test_sub_graph_error_on_from_node(valid_complex_graph):
-    ...
+    with pytest.raises(AssertionError) as e:
+        valid_complex_graph.sub_graph(from_node_id="NOT_EXISTING")
+        assert (
+            "Error sub graph from node NOT_EXISTING, node has to be in the execution graph"
+            == e
+        )
 
 
 def test_sub_graph_error_on_to_node(valid_complex_graph):
-    ...
+    with pytest.raises(AssertionError) as e:
+        valid_complex_graph.sub_graph(from_node_id="F", to_nodes_id=["NOT_EXIST"])
+        assert (
+            "Error sub graph to node NOT_EXIST, node has to be in the execution graph"
+            == e
+        )
+
+    with pytest.raises(AssertionError) as e:
+        valid_complex_graph.sub_graph(
+            from_node_id="F", to_nodes_id=["H", "M", "NOT_EXISTING"]
+        )
+        assert (
+            "Error sub graph to node NOT_EXISTING, node has to be in the execution graph"
+            == e
+        )