From 4cafaa0435a73ae70d9405b86c242c90c4c66033 Mon Sep 17 00:00:00 2001
From: Alberto Mercurio <alberto.mercurio96@gmail.com>
Date: Wed, 16 Oct 2024 00:11:15 +0200
Subject: [PATCH] Make separate function for allocations test

---
 test/basic.jl | 26 ++++++++++++++------------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/test/basic.jl b/test/basic.jl
index 2af3322..6956d7b 100644
--- a/test/basic.jl
+++ b/test/basic.jl
@@ -138,6 +138,11 @@ end
     @test ldiv!(op, u) ≈ (α * D) \ v
 end
 
+function apply_op!(H, du, u, p, t)
+    H(du, u, p, t)
+    return nothing
+end
+
 @testset "AddedOperator" begin
     A = rand(N, N) |> MatrixOperator
     B = rand(N, N) |> MatrixOperator
@@ -187,12 +192,9 @@ end
 
     # Allocations Tests
 
-    allocs_tot = @allocations mul!(v, op, u) # warmup
-    allocs_tot = @allocations mul!(v, op, u)
-    @test allocs_tot == 1 # BenchmarkTools.jl returns 0 instead of 1
-    allocs_tot = @allocations mul!(v, op, u, α, β) # warmup
-    allocs_tot = @allocations mul!(v, op, u, α, β)
-    @test allocs_tot == 1 # BenchmarkTools.jl returns 0 instead of 1
+    allocs_tot = @allocations apply_op!(op, v, u, (), 1.0) # warmup
+    allocs_tot = @allocations apply_op!(op, v, u, (), 1.0)
+    @test allocs_tot <= 1 # BenchmarkTools.jl returns 0 instead of 1
 
     ## Time-Dependent Coefficients
 
@@ -222,12 +224,12 @@ end
         p = (ω = 0.1,)
         t = 0.1
 
-        allocs_sparse = @allocations H_sparse(du, u, p, t) # warmup
-        allocs_sparse = @allocations H_sparse(du, u, p, t)
-        allocs_dense = @allocations H_dense(du, u, p, t) # warmup
-        allocs_dense = @allocations H_dense(du, u, p, t)
-        @test allocs_sparse == 1 # BenchmarkTools.jl returns 0 instead of 1
-        @test allocs_dense == 1 # BenchmarkTools.jl returns 0 instead of 1
+        allocs_sparse = @allocations apply_op!(H_sparse, du, u, p, t) # warmup
+        allocs_sparse = @allocations apply_op!(H_sparse, du, u, p, t)
+        allocs_dense = @allocations apply_op!(H_dense, du, u, p, t) # warmup
+        allocs_dense = @allocations apply_op!(H_dense, du, u, p, t)
+        @test allocs_sparse <= 1 # BenchmarkTools.jl returns 0 instead of 1
+        @test allocs_dense <= 1 # BenchmarkTools.jl returns 0 instead of 1
     end
 end