diff --git a/src/Linear/Linear.jl b/src/Linear/Linear.jl
index 334b6c9..dae1797 100644
--- a/src/Linear/Linear.jl
+++ b/src/Linear/Linear.jl
@@ -46,7 +46,7 @@ Track through a linear drift.
 - `beam` -- Input/output beam before/after tracking through
 - `ele`  -- `Linear.Drift` type element
 """
-function track!(beam::Beam, ele::Linear.Drift)
+function track!(beam::Beam, ele::Linear.Drift; work=get_work(beam, Val{0}()))
   L = ele.L
   v = beam.v
   
diff --git a/test/runtests.jl b/test/runtests.jl
index 4c90108..a1de870 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -5,89 +5,79 @@ using Test,
       BenchmarkTools,
       GTPSA
 
-@testset "BeamTracking.jl" begin
-    # Test setup
-    BenchmarkTools.DEFAULT_PARAMETERS.gctrial = false
-    BenchmarkTools.DEFAULT_PARAMETERS.evals = 2
-
-    # Linear tests -----------------------
-    d = Descriptor(6, 1) # 6 variables, order 1
-    beta_gamma_ref = 1.0
-    gamma_ref = sr_gamma(beta_gamma_ref)
-    species = Species("electron")
-    brho_ref = brho(massof(species), beta_gamma_ref, chargeof(species)) 
-    tol = 1e-14
-
-    # Linear drift test:
-    beam = Beam(d, beta_gamma_ref=beta_gamma_ref) # Creates Beam as identity map using GTPSA Descriptor
-    L_d = 5.0
-    drift = Linear.Drift(L=L_d)
-    track!(beam, drift)
-    
-    M_drift_expected = [1.0  L_d  0.0  0.0  0.0  0.0;
-                        0.0  1.0  0.0  0.0  0.0  0.0;
-                        0.0  0.0  1.0  L_d  0.0  0.0;
-                        0.0  0.0  0.0  1.0  0.0  0.0;
-                        0.0  0.0  0.0  0.0  1.0  L_d/gamma_ref^2;
-                        0.0  0.0  0.0  0.0  0.0  1.0]
+BenchmarkTools.DEFAULT_PARAMETERS.gctrial = false
+BenchmarkTools.DEFAULT_PARAMETERS.evals = 2
+  
+# Soon we generalize this to test_map...
+function test_matrix(ele, n_work, M_expected; type_stable=true, no_allocs=true, tol=1e-14, 
+                                              beta_gamma_ref=1.0, species=Species("electron")   )
     
-    p = Particle(beam, 1) # Converts SoA to 1 struct
+  d = Descriptor(6, 1) # 6 variables, order 1
+  beam = Beam(d, beta_gamma_ref=beta_gamma_ref, species=species) 
+  work = BeamTracking.get_work(beam, Val(n_work))
 
-    # 1) Correctness:
-    @test norm(GTPSA.jacobian(p.v) - M_drift_expected) < tol
-    # 2) Type-stability:
-    @test_opt track!(beam, drift) 
-    # 3) Allocations:
-    @test @ballocated(track!($beam, $drift)) == 0
+  track!(beam, ele, work=work)
+  p = Particle(beam, 1)
 
-    # Linear Quadrupole test:
-    # Focusing:
-    beam = Beam(d, beta_gamma_ref=beta_gamma_ref)
-    L_q = 1.2
-    K1n = 0.36
-    B1 = K1n*brho_ref
-    qf = Linear.Quadrupole(B1=B1,L=L_q)
-    work = BeamTracking.get_work(beam, Val{1}())
-    track!(beam, qf, work=work)
-
-    p = Particle(beam, 1) # Converts SoA to 1 struct
-    
-    M_qf_x = [cos(sqrt(K1n)*L_q)            sincu(sqrt(K1n)*L_q)*L_q;  
-              -sqrt(K1n)*sin(sqrt(K1n)*L_q) cos(sqrt(K1n)*L_q)      ;]
-    M_qf_y = [cosh(sqrt(K1n)*L_q)           sinhcu(sqrt(K1n)*L_q)*L_q; 
-              sqrt(K1n)*sinh(sqrt(K1n)*L_q) cosh(sqrt(K1n)*L_q)      ;]
-    M_qf_z = [1.0 L_q/gamma_ref^2;
-              0.0 1.0            ;]
+  # 1) Correctness
+  @test norm(GTPSA.jacobian(p.v) - M_expected) < tol 
+  # 2) Type stability
+  if type_stable
+    @test_opt track!(beam, ele, work=work)
+  end
+  # 3) No Allocations
+  if no_allocs
+    @test @ballocated(track!($beam, $ele, work=$work)) == 0 
+  end
+end
 
-    M_qf_expected = zeros(6,6)
-    M_qf_expected[1:2,1:2] = M_qf_x
-    M_qf_expected[3:4,3:4] = M_qf_y
-    M_qf_expected[5:6,5:6] = M_qf_z
-    # 1) Correctness:
-    @test norm(GTPSA.jacobian(p.v) - M_qf_expected) < tol
-    # 2) Type-stability:
-    @test_opt track!(beam, qf, work=work) 
-    # 3) Allocations:
-    @test @ballocated(track!($beam, $qf, work=$work)) == 0
+@testset "BeamTracking.jl" begin
+  # Linear tests -----------------------  
+  beta_gamma_ref = 1.0
+  gamma_ref = sr_gamma(beta_gamma_ref)
+  species = Species("electron")
+  brho_ref = brho(massof(species), beta_gamma_ref, chargeof(species)) 
+  tol = 1e-14 
 
-    # Defocusing:
-    beam = Beam(d, beta_gamma_ref=beta_gamma_ref)
-    qd = Linear.Quadrupole(B1=-B1,L=L_q)
-    track!(beam, qd, work=work)
+  # Linear drift test:
+  L_d = 5.0
+  drift = Linear.Drift(L=L_d)
+  M_drift_expected = [1.0  L_d  0.0  0.0  0.0  0.0;
+                      0.0  1.0  0.0  0.0  0.0  0.0;
+                      0.0  0.0  1.0  L_d  0.0  0.0;
+                      0.0  0.0  0.0  1.0  0.0  0.0;
+                      0.0  0.0  0.0  0.0  1.0  L_d/gamma_ref^2;
+                      0.0  0.0  0.0  0.0  0.0  1.0] 
+                      
+  test_matrix(drift, 0, M_drift_expected, beta_gamma_ref=beta_gamma_ref, species=species)
+  
+  # Linear Quadrupole test:
+  # Focusing:
+  L_q = 1.2
+  K1n = 0.36
+  B1 = K1n*brho_ref
+  qf = Linear.Quadrupole(B1=B1,L=L_q)
+  M_qf_x = [cos(sqrt(K1n)*L_q)            sincu(sqrt(K1n)*L_q)*L_q;  
+            -sqrt(K1n)*sin(sqrt(K1n)*L_q) cos(sqrt(K1n)*L_q)      ;]
+  M_qf_y = [cosh(sqrt(K1n)*L_q)           sinhcu(sqrt(K1n)*L_q)*L_q; 
+            sqrt(K1n)*sinh(sqrt(K1n)*L_q) cosh(sqrt(K1n)*L_q)      ;]
+  M_qf_z = [1.0 L_q/gamma_ref^2;
+            0.0 1.0            ;] 
+  M_qf_expected = zeros(6,6)
+  M_qf_expected[1:2,1:2] = M_qf_x
+  M_qf_expected[3:4,3:4] = M_qf_y
+  M_qf_expected[5:6,5:6] = M_qf_z
 
-    p = Particle(beam, 1)
+  test_matrix(qf, 1, M_qf_expected, beta_gamma_ref=beta_gamma_ref, species=species)
 
-    M_qd_expected = zeros(6,6)
-    M_qd_expected[1:2,1:2] = M_qf_y
-    M_qd_expected[3:4,3:4] = M_qf_x
-    M_qd_expected[5:6,5:6] = M_qf_z
+  # Defocusing:
+  qd = Linear.Quadrupole(B1=-B1,L=L_q)
+  M_qd_expected = zeros(6,6)
+  M_qd_expected[1:2,1:2] = M_qf_y
+  M_qd_expected[3:4,3:4] = M_qf_x
+  M_qd_expected[5:6,5:6] = M_qf_z 
 
-    # 1) Correctness:
-    @test norm(GTPSA.jacobian(p.v) - M_qd_expected) < tol
-    # 2) Type-stability:
-    @test_opt track!(beam, qd, work=work) 
-    # 3) Allocations:
-    @test @ballocated(track!($beam, $qd, work=$work)) == 0
+  test_matrix(qd, 1, M_qd_expected, beta_gamma_ref=beta_gamma_ref, species=species)
 end
 
 #include("linear.jl")