Fix async behaviour with patch contexts

Project.toml

@@ -4,14 +4,16 @@ keywords = ["testing", "mocking"]
 license = "MIT"
 desc = "Allows Julia function calls to be temporarily overloaded for purpose of testing"
 author = ["Curtis Vogt"]
-version = "0.7.3"
+version = "0.7.4"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+ContextVariablesX = "6add18c4-b38d-439d-96f6-d6bc489c04c5"
 ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
 
 [compat]
 Compat = "3.9"
+ContextVariablesX = "0.1.2"
 ExprTools = "0.1"
 julia = "1"
 

src/Mocking.jl

@@ -1,6 +1,7 @@
 module Mocking
 
 using Compat: mergewith
+using ContextVariablesX: @contextvar, with_context
 using ExprTools: splitdef, combinedef
 
 export @patch, @mock, Patch, apply

src/mock.jl

@@ -52,7 +52,7 @@ function get_alternate(pe::PatchEnv, target, args...)
     end
 end
 
-get_alternate(target, args...) = get_alternate(get_active_env(), target, args...)
+get_alternate(target, args...) = get_alternate(patch_env[], target, args...)
 
 function _debug_msg(method::Union{Method,Nothing}, target, args)
     call = "$target($(join(map(arg -> "::$(Core.Typeof(arg))", args), ", ")))"

src/patch.jl

@@ -126,19 +126,14 @@ end
 ```
 """
 function apply(body::Function, pe::PatchEnv)
-    prev_pe = get_active_env()
-    set_active_env(merge(prev_pe, pe))
-    try
-        return body()
-    finally
-        set_active_env(prev_pe)
-    end
+    merged_pe = merge(patch_env[], pe)
+    return with_context(body, patch_env => merged_pe)
 end
 
 function apply(body::Function, patches; debug::Bool=false)
     return apply(body, PatchEnv(patches, debug))
 end
 
-const PATCH_ENV = Ref{PatchEnv}(PatchEnv())
-set_active_env(pe::PatchEnv) = (PATCH_ENV[] = pe)
-get_active_env() = PATCH_ENV[]
+@contextvar patch_env = PatchEnv()
+set_active_env(body::Function, pe::PatchEnv) = with_context(body, patch_env => pe)
+get_active_env() = patch_env[]

test/async.jl

@@ -0,0 +1,40 @@
+@testset "tasks" begin
+    c = Condition()
+    ch = Channel{String}(1)
+    f() = "original"
+    function background()
+        # Wait until notified allowing us to control when this async code is executed
+        wait(c)
+        put!(ch, @mock f())
+        return nothing
+    end
+
+    p = @patch f() = "mocked"
+
+    @sync begin
+        # Task started outside patched context should not call patched functions.
+        @async background()
+        yield()
+
+        apply(p) do
+            @test (@mock f()) == "mocked"
+
+            notify(c)
+            @test take!(ch) == "original"
+
+            # Task started inside patched context should call patched functions.
+            @async background()
+            yield()
+            notify(c)
+            @test take!(ch) == "mocked"
+
+            # Task started inside patched context should call patched functions even when
+            # execution finishes outside of patched context.
+            @async background()
+            yield()
+        end
+
+        notify(c)
+        @test take!(ch) == "mocked"
+    end
+end

test/concept.jl

@@ -20,19 +20,19 @@
     for p in patches
         Mocking.apply!(pe, p)
     end
-    Mocking.set_active_env(pe)
 
-    @test (@mock multiply(2)) == 8        # calls mocked `multiply(::Int)`
-    @test (@mock multiply(0x2)) == 0x6    # calls mocked `multiply(::Integer)`
-    @test (@mock multiply(2//1)) == 4//1  # calls original `multiply(::Number)`
+    Mocking.set_active_env(pe) do
+        @test (@mock multiply(2)) == 8        # calls mocked `multiply(::Int)`
+        @test (@mock multiply(0x2)) == 0x6    # calls mocked `multiply(::Integer)`
+        @test (@mock multiply(2//1)) == 4//1  # calls original `multiply(::Number)`
 
-    @test (@mock multiply(2)) != multiply(2)
-    @test (@mock multiply(0x2)) != multiply(0x2)
-    @test (@mock multiply(2//1)) == multiply(2//1)
+        @test (@mock multiply(2)) != multiply(2)
+        @test (@mock multiply(0x2)) != multiply(0x2)
+        @test (@mock multiply(2//1)) == multiply(2//1)
+    end
 
     # Clean env
-    pe = Mocking.PatchEnv()
-    Mocking.set_active_env(pe)
+    @test Mocking.get_active_env() == Mocking.PatchEnv()
 
     # Ensure that original behaviour is restored
     @test (@mock multiply(2)) == 3

test/runtests.jl

@@ -28,4 +28,5 @@ using Mocking: anon_morespecific, anonymous_signature, dispatch, type_morespecif
     include("args.jl")
     include("merge.jl")
     include("nested_apply.jl")
+    include("async.jl")
 end