add downstream test

Signed-off-by: Thijs Baaijen <[email protected]>
PowerGridModel · Feb 5, 2025 · 2b99ea3 · 2b99ea3
2 parents afc8a17 + 5d541fb
commit 2b99ea3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
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
@@ -272,7 +272,9 @@ def get_downstream_nodes(self, node_id: int, start_node_ids: list[int], inclusiv
         """
         connected_node = self.find_first_connected(node_id, start_node_ids)
         path, _ = self.get_shortest_path(node_id, connected_node)
-        _, upstream_node, *_ = path  # path is at least 2 elements long or find_first_connected would have raised an error
+        _, upstream_node, *_ = (
+            path  # path is at least 2 elements long or find_first_connected would have raised an error
+        )
 
         return self.get_connected(node_id, [upstream_node], inclusive)
 

diff --git a/tests/unit/model/grids/test_grid_search.py b/tests/unit/model/grids/test_grid_search.py
@@ -26,14 +26,13 @@ def test_grid_get_nearest_substation_node_no_substation(basic_grid):
 
 
 class TestGetDownstreamNodes:
-
     def test_get_downstream_nodes(self):
-        grid = Grid.from_txt(["S1 11", "S1 2", "2 3", "3 5", "5 6", "2 4", "4 99", "99 100"])
+        grid = Grid.from_txt("S1 11", "S1 2", "2 3", "3 5", "5 6", "2 4", "4 99", "99 100")
         downstream_nodes = grid.get_downstream_nodes(node_id=3)
         assert [5, 6] == downstream_nodes
 
     def test_get_downstream_nodes_from_substation_node(self):
-        grid = Grid.from_txt(["S1 11", "S1 2", "2 3", "3 5", "5 6", "2 4", "4 99", "99 100"])
+        grid = Grid.from_txt("S1 11", "S1 2", "2 3", "3 5", "5 6", "2 4", "4 99", "99 100")
         with pytest.raises(NotImplementedError):
             grid.get_downstream_nodes(node_id=1)