From 98528a552b35b2ae1bdc17b8954399eb8ee31da8 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 17 Oct 2023 11:21:58 +0200
Subject: [PATCH] Add tests

---
 core/src/allocation.jl  | 17 +++++++----------
 core/test/allocation.jl | 32 ++++++++++++++++++++++++++++++++
 core/test/runtests.jl   |  1 +
 3 files changed, 40 insertions(+), 10 deletions(-)
 create mode 100644 core/test/allocation.jl

diff --git a/core/src/allocation.jl b/core/src/allocation.jl
index 0d59caa80..b7e4572e6 100644
--- a/core/src/allocation.jl
+++ b/core/src/allocation.jl
@@ -447,12 +447,7 @@ function get_allocation_model(
     )
 end
 
-function allocate!(
-    p::Parameters,
-    allocation_model::AllocationModel,
-    t::Float64;
-    return_flows::Bool = false,
-)::Union{Nothing, Vector{Float64}}
+function allocate!(p::Parameters, allocation_model::AllocationModel, t::Float64;)::Nothing
     (; node_id_mapping, source_edge_mapping, model) = allocation_model
     (; user, connectivity) = p
     (; priorities, demand) = user
@@ -498,9 +493,11 @@ function allocate!(
 
     optimize!(model)
 
-    if return_flows
-        return value.(model[:F])
-    else
-        return nothing
+    for (subnetwork_node_id, (AG_node_id, AG_node_type)) in node_id_mapping
+        if AG_node_type == :user
+            user_idx = findsorted(user.node_id, subnetwork_node_id)
+            base_name = "A_user_$AG_node_id"
+            user.allocated[user_idx] .= value.(model[Symbol(base_name)])
+        end
     end
 end
diff --git a/core/test/allocation.jl b/core/test/allocation.jl
new file mode 100644
index 000000000..122318df8
--- /dev/null
+++ b/core/test/allocation.jl
@@ -0,0 +1,32 @@
+import Ribasim
+
+@testset "Allocation solve" begin
+    toml_path = normpath(@__DIR__, "../../generated_testmodels/subnetwork/subnetwork.toml")
+    @test ispath(toml_path)
+    cfg = Ribasim.Config(toml_path)
+    gpkg_path = Ribasim.input_path(cfg, cfg.geopackage)
+    db = SQLite.DB(gpkg_path)
+
+    p = Ribasim.Parameters(db, cfg)
+    close(db)
+
+    # Inputs specific for this test model
+    subgraph_node_ids = unique(keys(p.lookup))
+    source_edge_ids = [p.connectivity.edge_ids_flow[(1, 2)]]
+    flow = Ribasim.get_tmp(p.connectivity.flow, 0)
+    flow[1, 2] = 4.5 # Source flow
+    Δt_allocation = 24.0 * 60^2
+    t = 0.0
+
+    allocation_model =
+        Ribasim.get_allocation_model(p, subgraph_node_ids, source_edge_ids, Δt_allocation)
+    Ribasim.allocate!(p, allocation_model, t)
+
+    F = value.(allocation_model.model[:F])
+    @test F ≈ [3.0, 3.0, 0.5, 3.5, 1.0, 4.5]
+
+    allocated = p.user.allocated
+    @test allocated[1] ≈ [0.0, 1.0, 0.0]
+    @test allocated[2] ≈ [0.0, 0.5, 0.0]
+    @test allocated[3] ≈ [3.0, 0.0, 0.0]
+end
diff --git a/core/test/runtests.jl b/core/test/runtests.jl
index 8a73f26bc..e143dc49d 100644
--- a/core/test/runtests.jl
+++ b/core/test/runtests.jl
@@ -12,6 +12,7 @@ using SafeTestsets: @safetestset
     @safetestset "Basic Model Interface" include("bmi.jl")
     @safetestset "Utility functions" include("utils.jl")
     @safetestset "Control" include("control.jl")
+    @safetesteset "Allocation" include("allocation.jl")
     @safetestset "Time" include("time.jl")
     @safetestset "Docs" include("docs.jl")
     @safetestset "Command Line Interface" include("cli.jl")