From c975b0c0bac347fdf66398740fa7b5bc11950ddc Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 11 Apr 2024 11:19:54 +0200
Subject: [PATCH] Fix table sorting

---
 core/src/validation.jl           |  5 +++++
 python/ribasim/ribasim/config.py | 10 +++++++++-
 python/ribasim/tests/test_io.py  |  1 +
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/core/src/validation.jl b/core/src/validation.jl
index 1280049dc..d4b46b203 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -102,7 +102,10 @@ sort_by_id_state_level(row) = (row.node_id, row.control_state, row.level)
 sort_by_priority(row) = (row.node_id, row.priority)
 sort_by_priority_time(row) = (row.node_id, row.priority, row.time)
 sort_by_subgrid_level(row) = (row.subgrid_id, row.basin_level)
+sort_by_condition(row) =
+    (row.node_id, row.listen_node_type, row.listen_node_id, row.variable, row.greater_than)
 sort_by_id_name(row) = (row.node_id, row.name)
+sort_by_truth_state(row) = (row.node_id, row.truth_state)
 
 # get the right sort by function given the Schema, with sort_by_id as the default
 sort_by_function(table::StructVector{<:Legolas.AbstractRecord}) = sort_by_id
@@ -112,6 +115,8 @@ sort_by_function(table::StructVector{UserDemandStaticV1}) = sort_by_priority
 sort_by_function(table::StructVector{UserDemandTimeV1}) = sort_by_priority_time
 sort_by_function(table::StructVector{BasinSubgridV1}) = sort_by_subgrid_level
 sort_by_function(table::StructVector{DiscreteControlCompoundvariableV1}) = sort_by_id_name
+sort_by_function(table::StructVector{DiscreteControlConditionV1}) = sort_by_condition
+sort_by_function(table::StructVector{DiscreteControlLogicV1}) = sort_by_truth_state
 
 const TimeSchemas = Union{
     BasinTimeV1,
diff --git a/python/ribasim/ribasim/config.py b/python/ribasim/ribasim/config.py
index 3fed1f8a7..e00f20ed2 100644
--- a/python/ribasim/ribasim/config.py
+++ b/python/ribasim/ribasim/config.py
@@ -294,7 +294,15 @@ class DiscreteControl(MultiNodeModel):
     )
     condition: TableModel[DiscreteControlConditionSchema] = Field(
         default_factory=TableModel[DiscreteControlConditionSchema],
-        json_schema_extra={"sort_keys": ["node_id"]},
+        json_schema_extra={
+            "sort_keys": [
+                "node_id",
+                "listen_node_type",
+                "listen_node_id",
+                "variable",
+                "greater_than",
+            ]
+        },
     )
     logic: TableModel[DiscreteControlLogicSchema] = Field(
         default_factory=TableModel[DiscreteControlLogicSchema],
diff --git a/python/ribasim/tests/test_io.py b/python/ribasim/tests/test_io.py
index bd6fcc7a8..e9b724eaa 100644
--- a/python/ribasim/tests/test_io.py
+++ b/python/ribasim/tests/test_io.py
@@ -102,6 +102,7 @@ def test_sort(level_setpoint_with_minmax, tmp_path):
     assert table.df.iloc[0]["greater_than"] == 15.0
     assert table._sort_keys == [
         "node_id",
+        "listen_node_type",
         "listen_node_id",
         "variable",
         "greater_than",