nrow

Deltares · Jan 26, 2024 · 89e2ca9 · 89e2ca9
1 parent 6301379
commit 89e2ca9
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/python/ribasim/tests/test_model.py b/python/ribasim/tests/test_model.py
@@ -137,11 +137,15 @@ def test_plot(discrete_control_of_pid_control):
 def test_write_adds_fid_in_tables(basic, tmp_path):
     model_orig = basic
     # for node an explicit index was provided
+    nrow = len(model_orig.network.node.df)
     assert model_orig.network.node.df.index.name == "fid"
-    assert model_orig.network.node.df.index.equals(pd.RangeIndex(start=1, stop=18))
+    assert model_orig.network.node.df.index.equals(
+        pd.RangeIndex(start=1, stop=nrow + 1)
+    )
     # for edge no index was provided, but it still needs to write it to file
+    nrow = len(model_orig.network.edge.df)
     assert model_orig.network.edge.df.index.name is None
-    assert model_orig.network.edge.df.index.equals(pd.RangeIndex(start=0, stop=17))
+    assert model_orig.network.edge.df.index.equals(pd.RangeIndex(start=0, stop=nrow))
 
     model_orig.write(tmp_path / "basic/ribasim.toml")
     with connect(tmp_path / "basic/database.gpkg") as connection: