From 38016b6c72d28d3586cb922fce2006dc4c35f847 Mon Sep 17 00:00:00 2001
From: Albert de Montserrat <albert.de.montserrat@gmail.com>
Date: Mon, 18 Nov 2024 14:52:40 +0100
Subject: [PATCH] Mask test

---
 test/test_mask.jl | 63 +++++++++++++++++++++++++----------------------
 1 file changed, 34 insertions(+), 29 deletions(-)

diff --git a/test/test_mask.jl b/test/test_mask.jl
index 9eb7e2d0..ecd2f056 100644
--- a/test/test_mask.jl
+++ b/test/test_mask.jl
@@ -1,41 +1,46 @@
 using Test
+using JustRelax, JustRelax.JustRelax2D
+import JustRelax: Mask, apply_mask, apply_mask!
+import JustRelax.JustRelax2D as JR2
 
-ni = 10, 10
+@testset "Mask" begin
+    ni = 10, 10
 
-# Test basics 
-m  = Mask(ni...)
+    # Test basics 
+    m  = Mask(ni...)
 
-@test size(m)      == ni
-@test length(m)    == prod(ni)
-@test axes(m)      == Base.OneTo(ni), Base.OneTo(ni)
-@test eachindex(m) == Base.OneTo(prod(ni))
-@test similar(m)   isa Mask{Matrix{Float64}}
-@test !all(m) 
+    @test size(m)      == ni
+    @test length(m)    == prod(ni)
+    @test axes(m)      == (Base.OneTo(ni[1]), Base.OneTo(ni[2]))
+    @test eachindex(m) == Base.OneTo(prod(ni))
+    @test similar(m)   isa Mask{Matrix{Float64}}
+    @test !all(m) 
 
-m  = Mask(ni..., 4:7, 4:7)
-@test all(isone, m[4:7, 4:7])
+    m  = JR2.Mask(ni..., 4:7, 4:7)
+    @test all(isone, m[4:7, 4:7])
 
-# test masking
-A  = rand(ni...)
-B  = zeros(ni...)
-B[4:7, 4:7] .= 5
-m  = Mask(ni..., 4:7, 4:7)
+    # test masking
+    A  = rand(ni...)
+    B  = zeros(ni...)
+    B[4:7, 4:7] .= 5
+    m  = JR2.Mask(ni..., 4:7, 4:7)
 
-C = apply_mask(A, B, m)
-@test all(C[4:7, 4:7] .== 5)
+    C = apply_mask(A, B, m)
+    @test all(C[4:7, 4:7] .== 5)
 
-apply_mask!(A, B, m)
-@test all(A[4:7, 4:7] .== 5)
+    apply_mask!(A, B, m)
+    @test all(A[4:7, 4:7] .== 5)
 
-A  = rand(ni...)
-@test apply_mask(A, B, m, 1, 1) == A[1, 1]
-@test apply_mask(A, B, m, 5, 5) == 5
+    A  = rand(ni...)
+    @test apply_mask(A, B, m, 1, 1) == A[1, 1]
+    @test apply_mask(A, B, m, 5, 5) == 5
 
-apply_mask!(A, B, m, 1, 1)
-apply_mask!(A, B, m, 5, 5)
+    apply_mask!(A, B, m, 1, 1)
+    apply_mask!(A, B, m, 5, 5)
 
-@test A[1, 1] != 5
-@test A[5, 5] == 5
+    @test A[1, 1] != 5
+    @test A[5, 5] == 5
 
-@test isone(inv(m, 1, 1))
-@test iszero(inv(m, 5, 5))
\ No newline at end of file
+    @test isone(inv(m, 1, 1))
+    @test iszero(inv(m, 5, 5))
+end
\ No newline at end of file