diff --git a/test/PDE_UDE_solve.jl b/test/PDE_UDE_solve.jl
index a35b95f..a07e0e3 100644
--- a/test/PDE_UDE_solve.jl
+++ b/test/PDE_UDE_solve.jl
@@ -1,6 +1,7 @@
 
 
-function pde_solve_test(atol::F, save_refs::Bool = false; MB::Bool=false, fast::Bool=true) where {F <: AbstractFloat}
+function pde_solve_test(; rtol::F, atol::F, save_refs::Bool=false, MB::Bool=false, fast::Bool=true) where {F <: AbstractFloat}
+
     println("PDE solving with MB = $MB")
     working_dir = joinpath(homedir(), "OGGM/Huginn_tests")
     
@@ -9,9 +10,10 @@ function pde_solve_test(atol::F, save_refs::Bool = false; MB::Bool=false, fast::
                                               workers=2,
                                               ice_thickness_source = "Farinotti19"),
                         simulation = SimulationParameters(use_MB=MB,
-                                                        velocities=false,
-                                                        tspan=(2010.0, 2015.0),
-                                                        working_dir = Huginn.root_dir)
+                                                          velocities=false,
+                                                          tspan=(2010.0, 2015.0),
+                                                          working_dir = Huginn.root_dir),
+                        solver = SolverParameters(reltol=1e-12)
                         ) 
 
     ## Retrieving gdirs and climate for the following glaciers
@@ -24,13 +26,6 @@ function pde_solve_test(atol::F, save_refs::Bool = false; MB::Bool=false, fast::
         "RGI60-07.00274", "RGI60-07.01323",  "RGI60-01.17316"]
     end
 
-    # Load reference values for the simulation
-    if MB
-        PDE_refs = load(joinpath(Huginn.root_dir, "test/data/PDE/PDE_refs_MB.jld2"))["results"]
-    else
-        PDE_refs = load(joinpath(Huginn.root_dir, "test/data/PDE/PDE_refs_noMB.jld2"))["results"]
-    end
-
     model = Model(iceflow = SIA2Dmodel(params))
 
     # We retrieve some glaciers for the simulation
@@ -51,6 +46,13 @@ function pde_solve_test(atol::F, save_refs::Bool = false; MB::Bool=false, fast::
         end
     end
     
+    # Load reference values for the simulation
+    if MB
+        PDE_refs = load(joinpath(Huginn.root_dir, "test/data/PDE/PDE_refs_MB.jld2"))["results"]
+    else
+        PDE_refs = load(joinpath(Huginn.root_dir, "test/data/PDE/PDE_refs_noMB.jld2"))["results"]
+    end
+
     # # Run one epoch of the UDE training
     # θ = zeros(10) # dummy data for the NN
     # UA_f = zeros(10)
@@ -59,6 +61,7 @@ function pde_solve_test(atol::F, save_refs::Bool = false; MB::Bool=false, fast::
     # H_V_preds = @time predict_iceflow(θ, UA_f, gdirs, context_batches, UDE_settings, mb_model) # Array{(H_pred, V̄x_pred, V̄y_pred, rgi_id)}
 
     let results=prediction.results
+
     for result in results
         let result=result, test_ref=nothing, UDE_pred=nothing
         for PDE_ref in PDE_refs
@@ -79,19 +82,22 @@ function pde_solve_test(atol::F, save_refs::Bool = false; MB::Bool=false, fast::
             mkdir(test_plot_path)
         end
         MB ? vtol = 30.0*atol : vtol = 12.0*atol # a little extra tolerance for UDE surface velocities
+
         ### PDE ###
-        plot_test_error(result, test_ref, "H", result.rgi_id, atol, MB)
+        plot_test_error(result, test_ref, "H",  result.rgi_id, atol, MB)
         plot_test_error(result, test_ref, "Vx", result.rgi_id, vtol, MB)
         plot_test_error(result, test_ref, "Vy", result.rgi_id, vtol, MB)
+        
         ### UDE ###
         # Huginn.plot_test_error(UDE_pred, test_ref, "H", rgi_id, atol, MB)
         # Huginn.plot_test_error(UDE_pred, test_ref, "Vx", rgi_id, vtol, MB)
         # Huginn.plot_test_error(UDE_pred, test_ref, "Vy", rgi_id, vtol, MB)
 
         # Test that the PDE simulations are correct
-        @test all(isapprox.(result.H[end], test_ref.H[end], atol=atol))
-        @test all(isapprox.(result.Vx, test_ref.Vx, atol=vtol))
-        @test all(isapprox.(result.Vy, test_ref.Vy, atol=vtol))
+        @test all(isapprox.(result.H[end], test_ref.H[end], rtol=rtol, atol=atol))
+        @test all(isapprox.(result.Vx, test_ref.Vx, rtol=rtol, atol=vtol))
+        @test all(isapprox.(result.Vy, test_ref.Vy, rtol=rtol, atol=vtol))
+        
         # Test that the UDE simulations are correct
         # @test all(isapprox.(UDE_pred[1], test_ref["H"], atol=atol))
         # @test all(isapprox.(UDE_pred[2], test_ref["Vx"], atol=vtol)) 
diff --git a/test/data/PDE/PDE_refs_noMB.jld2 b/test/data/PDE/PDE_refs_noMB.jld2
index 7ce09c1..69a1e8b 100644
Binary files a/test/data/PDE/PDE_refs_noMB.jld2 and b/test/data/PDE/PDE_refs_noMB.jld2 differ
diff --git a/test/mass_conservation.jl b/test/mass_conservation.jl
index fbf885e..93affe0 100644
--- a/test/mass_conservation.jl
+++ b/test/mass_conservation.jl
@@ -21,7 +21,7 @@ Arguments
     - `rtol`: Relative tolerance
     - `save_plot`: Optional bool to save plot during simulation
 """
-function unit_mass_test(; H₀, B, A, t_sim, Δx, Δy, rtol=0.02, save_plot=false)
+function unit_mass_test(; H₀, B, A, n, t_sim, Δx, Δy, rtol=0.02, save_plot=false)
 
     # Get parameters for a simulation 
     parameters = Parameters(simulation=SimulationParameters(tspan=(0.0, t_sim),
@@ -29,7 +29,7 @@ function unit_mass_test(; H₀, B, A, t_sim, Δx, Δy, rtol=0.02, save_plot=fals
                                                             use_iceflow=true,
                                                             multiprocessing=true,
                                                             workers=1),
-                            physical=PhysicalParameters(A=A),
+                            physical=PhysicalParameters(),
                             solver=SolverParameters(reltol=1e-12))
 
     model = Model(iceflow = SIA2Dmodel(parameters))
@@ -39,8 +39,8 @@ function unit_mass_test(; H₀, B, A, t_sim, Δx, Δy, rtol=0.02, save_plot=fals
     nx, ny = size(H₀)
 
     # Define glacier object
-    glacier = Glacier(rgi_id = "toy", H₀ = H₀, S = S, B = B, 
-                      Δx=Δx, Δy=Δy, nx=nx, ny=ny)
+    glacier = Glacier2D(rgi_id = "toy", H₀ = H₀, S = S, B = B, A = A, n = n, 
+                        Δx=Δx, Δy=Δy, nx=nx, ny=ny)
     glaciers = Vector{Sleipnir.AbstractGlacier}([glacier])
 
     prediction = Prediction(model, glaciers, parameters)
@@ -99,7 +99,7 @@ function unit_mass_flatbed_test(; rtol)
                     H₀ = zeros((nx,ny))
                     @views H₀[floor(Int,nx/3):floor(Int,2nx/3), floor(Int,ny/3):floor(Int,2ny/3)] .= 400
                 end
-                unit_mass_test(; H₀=H₀, B=B, A=A, t_sim=10.0, Δx=50.0, Δy=50.0, rtol=rtol, save_plot=false)
+                unit_mass_test(; H₀=H₀, B=B, A=A, n=3.0, t_sim=10.0, Δx=50.0, Δy=50.0, rtol=rtol, save_plot=false)
             end
         end
     end
diff --git a/test/runtests.jl b/test/runtests.jl
index 89dc978..9d3802d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -23,10 +23,9 @@ ENV["GKSwstype"]="nul"
 
 @testset "Solver parameters construction with default variables" params_constructor_default()
 
-@testset "Halfar Solution" halfar_test(rtol=0.02, atol=1.0)
+@testset "Halfar Solution" halfar_test(; rtol=0.02, atol=1.0)
 
-atol = 0.01
-@testset "PDE solving integration tests" pde_solve_test(atol; MB=false, fast=true)
+@testset "PDE solving integration tests" pde_solve_test(; rtol=0.01, atol=0.01, save_refs=false, MB=false, fast=true)
 
-@testset "Conservation of Mass - Flat Bed" unit_mass_flatbed_test(rtol=1.0e-7)
+@testset "Conservation of Mass - Flat Bed" unit_mass_flatbed_test(; rtol=1.0e-7)