diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py index 2a4a577..2ce1b88 100644 --- a/src/power_grid_model_ds/_core/model/graphs/models/base.py +++ b/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) diff --git a/tests/unit/model/graphs/test_container.py b/tests/unit/model/graphs/test_container.py index ada2fb0..a068573 100644 --- a/tests/unit/model/graphs/test_container.py +++ b/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]