diff --git a/src/visualization.jl b/src/visualization.jl
index 7918918c..f4c44d39 100644
--- a/src/visualization.jl
+++ b/src/visualization.jl
@@ -186,10 +186,10 @@ end
 
         for (i, (name, integral)) in enumerate(pairs(ints))
             name in exclude && continue
-            name_function = getfield(@__MODULE__, name)
+            quantity = cb.affect!.analysis_integrals[i]
             @series begin
                 subplot --> subplot
-                label := pretty_form_utf(name_function) * " " * label_extension
+                label := pretty_form_utf(quantity) * " " * label_extension
                 title --> "change of invariants"
                 xguide --> "t"
                 yguide --> "change of invariants"
diff --git a/test/test_visualization.jl b/test/test_visualization.jl
index 56d6ff27..f7e1a9be 100644
--- a/test/test_visualization.jl
+++ b/test/test_visualization.jl
@@ -5,7 +5,10 @@ using DispersiveShallowWater
 using Plots
 
 @testset "Visualization" begin
-    trixi_include(@__MODULE__, default_example(), tspan = (0.0, 1.0))
+    custom_integral(q, equations) = q[1]^2
+    DispersiveShallowWater.pretty_form_utf(::typeof(custom_integral)) = "∫η²"
+    trixi_include(@__MODULE__, default_example(), tspan = (0.0, 1.0),
+                  extra_analysis_integrals = (waterheight_total, custom_integral))
     @test_nowarn plot(semi => sol)
     @test_nowarn plot!(semi => sol, plot_initial = true)
     @test_nowarn plot(semi, sol, conversion = prim2cons, plot_bathymetry = false)