From 37395d97f9ce803e8af6aa5790e94d28c660f960 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 26 Mar 2024 14:38:43 +0100
Subject: [PATCH] Add tests

---
 core/src/callback.jl   |  2 +-
 core/test/time_test.jl | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 8d1b02278..a6abacdcd 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -158,7 +158,7 @@ function save_vertical_flux(u, t, integrator)
     Δt = get_Δt(integrator, graph)
     vertical_flux_mean = copy(vertical_flux_integrated)
     vertical_flux_mean ./= Δt
-    fill!(vertical_flux_mean, 0.0)
+    fill!(vertical_flux_integrated, 0.0)
 
     return vertical_flux_mean
 end
diff --git a/core/test/time_test.jl b/core/test/time_test.jl
index 947d56908..7f8862a19 100644
--- a/core/test/time_test.jl
+++ b/core/test/time_test.jl
@@ -22,3 +22,45 @@
     # some difference is expected since the modeled flow is for the period up to t
     @test isapprox(flow_1_to_2.flow_rate, flow_expected, rtol = 0.005)
 end
+
+@test "vertical_flux_means" begin
+    using DataFrames: DataFrame
+
+    toml_path =
+        normpath(@__DIR__, "../../generated_testmodels/basic_transient/ribasim.toml")
+    @test ispath(toml_path)
+    model = Ribasim.run(toml_path)
+    basin = model.integrator.p.basin
+    basin_table = DataFrame(Ribasim.basin_table(model))
+    filter!(:node_id => id -> id == 1, basin_table)
+    basin_table[!, "time"] .=
+        Ribasim.seconds_since.(basin_table.time, model.config.starttime)
+
+    # No vertical flux data for last saveat
+    t_end = last(basin_table).time
+    data_end = filter(:time => t -> t == t_end, basin_table)
+    @test all(data_end.precipitation .== 0)
+    @test all(data_end.evaporation .== 0)
+    @test all(data_end.drainage .== 0)
+    @test all(data_end.infiltration .== 0)
+
+    time_table = DataFrame(basin.time)
+    filter!(:node_id => id -> id == 1, time_table)
+    time_table[!, "time"] .= Ribasim.seconds_since.(time_table.time, model.config.starttime)
+    fixed_area = basin.area[1][end]
+    area = [
+        Ribasim.get_area_and_level(basin, 1, storage)[1] for storage in basin_table.storage
+    ]
+    area = (area[2:end] .+ area[1:(end - 1)]) ./ 2
+    @test all(
+        basin_table.precipitation[1:(end - 1)] .≈
+        fixed_area .* time_table.precipitation[1:(end - 1)],
+    )
+    @test all(
+        isapprox(
+            basin_table.evaporation[1:(end - 1)],
+            area .* time_table.potential_evaporation[1:(end - 1)];
+            rtol = 1e-4,
+        ),
+    )
+end