fix: delete_branch3 used wrong argument name (#29)

Signed-off-by: Vincent Koppen <[email protected]>

Author: vincentkoppen

src/power_grid_model_ds/_core/model/graphs/models/base.py

@@ -158,9 +158,9 @@ def delete_branch_array(self, branch_array: BranchArray, raise_on_fail: bool = T
             if self._branch_is_relevant(branch):
                 self.delete_branch(branch.from_node.item(), branch.to_node.item(), raise_on_fail=raise_on_fail)
 
-    def delete_branch3_array(self, branch_array: Branch3Array, raise_on_fail: bool = True) -> None:
+    def delete_branch3_array(self, branch3_array: Branch3Array, raise_on_fail: bool = True) -> None:
         """Delete all branch3s in the branch3 array from the graph."""
-        for branch3 in branch_array:
+        for branch3 in branch3_array:
             branches = _get_branch3_branches(branch3)
             self.delete_branch_array(branches, raise_on_fail=raise_on_fail)
 

tests/unit/model/graphs/test_container.py

@@ -26,6 +26,57 @@ def test_from_arrays(basic_grid):
     assert set(basic_grid.node.id) == set(graphs.complete_graph.external_ids)
 
 
+@pytest.fixture
+def graph_container_with_5_nodes():
+    graph_container = GraphContainer.empty()
+    for node_id in range(1, 6):
+        node = NodeArray.empty(1)
+        node.id = node_id
+        graph_container.add_node(node)
+    return graph_container
+
+
+@pytest.fixture
+def three_winding_transformers():
+    three_winding_transformers = ThreeWindingTransformerArray.empty(2)
+    three_winding_transformers.id = [301, 302]
+    three_winding_transformers.node_1 = [1, 1]
+    three_winding_transformers.node_2 = [2, 4]
+    three_winding_transformers.node_3 = [3, 5]
+    three_winding_transformers.status_1 = [1, 1]
+    three_winding_transformers.status_2 = [1, 1]
+    three_winding_transformers.status_3 = [0, 1]
+
+    return three_winding_transformers
+
+
+def test_add_branch3(graph_container_with_5_nodes, three_winding_transformers):
+    graph_container_with_5_nodes.add_branch3(three_winding_transformers)
+    for from_node, to_node in [(1, 2), (1, 4), (1, 5), (4, 5)]:
+        assert graph_container_with_5_nodes.active_graph.has_branch(from_node, to_node)
+        assert graph_container_with_5_nodes.complete_graph.has_branch(from_node, to_node)
+
+    for from_node, to_node in [(1, 3), (2, 3)]:
+        assert not graph_container_with_5_nodes.active_graph.has_branch(from_node, to_node)
+        assert graph_container_with_5_nodes.complete_graph.has_branch(from_node, to_node)
+
+
+def test_delete_branch3(graph_container_with_5_nodes, three_winding_transformers):
+    graph_container_with_5_nodes.add_branch3(three_winding_transformers)
+    graph_container_with_5_nodes.delete_branch3(three_winding_transformers[0])
+
+    assert not graph_container_with_5_nodes.active_graph.has_branch(1, 2)
+    assert not graph_container_with_5_nodes.complete_graph.has_branch(1, 2)
+    for from_node, to_node in [(1, 4), (1, 5), (4, 5)]:
+        assert graph_container_with_5_nodes.active_graph.has_branch(from_node, to_node)
+        assert graph_container_with_5_nodes.complete_graph.has_branch(from_node, to_node)
+    graph_container_with_5_nodes.delete_branch3(three_winding_transformers[1])
+
+    for from_node, to_node in [(1, 2), (1, 4), (1, 5), (4, 5)]:
+        assert not graph_container_with_5_nodes.active_graph.has_branch(from_node, to_node)
+        assert not graph_container_with_5_nodes.complete_graph.has_branch(from_node, to_node)
+
+
 def test_from_arrays_active_three_winding(basic_grid):
     nodes = NodeArray.zeros(3)
     nodes.id = [1000, 1001, 1002]