diff --git a/test/test_enzyme.jl b/test/test_enzyme.jl
index 45faab48ff..cc86656889 100644
--- a/test/test_enzyme.jl
+++ b/test/test_enzyme.jl
@@ -157,7 +157,8 @@ end
     end
 end
 
-@testset "Enzyme on advection and diffusion" begin
+
+@testset "Enzyme for advection and diffusion with various boundary conditions" begin
     Nx = Ny = 64
     Nz = 8
 
@@ -183,16 +184,17 @@ end
     fill_halo_regions!(u)
     fill_halo_regions!(v)
 
-    @inline function tracer_flux(x, y, t, c, p)
+    @inline function tracer_flux(i, j, grid, clock, model_fields, p)
         c₀ = p.surface_tracer_concentration
         u★ = p.piston_velocity
-        return - u★ * (c₀ - c)
+        return - u★ * (c₀ - model_fields.c[i, j, p.level])
     end
 
     parameters = (surface_tracer_concentration = 1,
-                  piston_velocity = 0.1)
+                  piston_velocity = 0.1,
+                  level = Nz)
 
-    top_c_bc = FluxBoundaryCondition(tracer_flux, field_dependencies=:c; parameters)
+    top_c_bc = FluxBoundaryCondition(tracer_flux; discrete_form=true, parameters)
     c_bcs = FieldBoundaryConditions(top=top_c_bc)
 
     # TODO:
@@ -200,37 +202,51 @@ end
     # 2. Add surface fluxes
     # 3. Do a problem where we invert for the tracer fluxes (maybe with CATKE)
 
-    model = HydrostaticFreeSurfaceModel(; grid,
-                                        tracer_advection = WENO(),
-                                        tracers = :c,
-                                        velocities = PrescribedVelocityFields(; u, v),
-                                        closure = diffusion)
-
-    # Compute derivative by hand
-    κ₁, κ₂ = 0.9, 1.1
-    c²₁ = stable_diffusion!(model, 1, κ₁)
-    c²₂ = stable_diffusion!(model, 1, κ₂)
-    dc²_dκ_fd = (c²₂ - c²₁) / (κ₂ - κ₁)
-
-    # Now for real
-    amplitude = 1.0
-    κ = 1.0
-    dmodel = Enzyme.make_zero(model)
-    set_diffusivity!(dmodel, 0)
-
-    dc²_dκ = autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse),
-                      stable_diffusion!,
-                      Duplicated(model, dmodel),
-                      Const(amplitude),
-                      Active(κ))
-
-    @info """ \n
+    model_no_bc = HydrostaticFreeSurfaceModel(; grid,
+                                              tracer_advection = WENO(),
+                                              tracers = :c,
+                                              velocities = PrescribedVelocityFields(; u, v),
+                                              closure = diffusion)
+
+    model_bc = HydrostaticFreeSurfaceModel(; grid,
+                                           tracer_advection = WENO(),
+                                           tracers = :c,
+                                           velocities = PrescribedVelocityFields(; u, v),
+                                           boundary_conditions = (; c=c_bcs),
+                                           closure = diffusion)
+
+    models = [model_no_bc, model_bc]
+
+    @show "Advection-diffusion results, first without then with flux BC"
+
+    for i in 1:2
+        # Compute derivative by hand
+        κ₁, κ₂ = 0.9, 1.1
+        c²₁ = stable_diffusion!(models[i], 1, κ₁)
+        c²₂ = stable_diffusion!(models[i], 1, κ₂)
+        dc²_dκ_fd = (c²₂ - c²₁) / (κ₂ - κ₁)
+
+        # Now for real
+        amplitude = 1.0
+        κ = 1.0
+        dmodel = Enzyme.make_zero(models[i])
+        set_diffusivity!(dmodel, 0)
+
+        dc²_dκ = autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse),
+                          stable_diffusion!,
+                          Duplicated(models[i], dmodel),
+                          Const(amplitude),
+                          Active(κ))
+
+        @info """ \n
+        Advection-diffusion:
         Enzyme computed $dc²_dκ
         Finite differences computed $dc²_dκ_fd
-    """
+        """
 
-    tol = 0.01
-    rel_error = abs(dc²_dκ[1][3] - dc²_dκ_fd) / abs(dc²_dκ_fd)
-    @test rel_error < tol
+        tol = 0.01
+        rel_error = abs(dc²_dκ[1][3] - dc²_dκ_fd) / abs(dc²_dκ_fd)
+        @test rel_error < tol
+    end
 end