Use ScopedValue to fix async issue (#112)

* Use ContextVariablesX to fix async issue

* Use ScopedValues rather than ContextVariables

* Apply suggestions from code review

Co-authored-by: Curtis Vogt <curtis.vogt@gmail.com>

* Define missing PATCH_ENV

* Refactor tasks testset into async scope

* Import ScopedValue alongside other imports

* Use VERSION check instead

---------

Co-authored-by: Curtis Vogt <curtis.vogt@gmail.com>

src/Mocking.jl

@@ -3,6 +3,13 @@ module Mocking
 using Compat: mergewith
 using ExprTools: splitdef, combinedef
 
+# Available in Julia 1.11+: https://github.com/JuliaLang/julia/pull/50958
+# We cannot use ScopedValues.jl for backwards compatability as that implementation breaks
+# `@test_logs`.
+if VERSION >= v"1.11.0-DEV.482"
+    using Base: ScopedValue, with
+end
+
 export @patch, @mock, Patch, apply
 
 include("expr.jl")

src/mock.jl

@@ -49,7 +49,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,28 @@ 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_active_env(body, 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[]
+# https://github.com/JuliaLang/julia/pull/50958
+if VERSION >= v"1.11.0-DEV.482"
+    const PATCH_ENV = ScopedValue(PatchEnv())
+    with_active_env(body::Function, pe::PatchEnv) = with(body, PATCH_ENV => pe)
+else
+    const PATCH_ENV = Ref{PatchEnv}(PatchEnv())
+
+    function with_active_env(body::Function, pe::PatchEnv)
+        old_pe = PATCH_ENV[]
+        try
+            PATCH_ENV[] = pe
+            body()
+        finally
+            PATCH_ENV[] = old_pe
+        end
+    end
+end

test/async.jl → test/async-scope.jl

@@ -1,4 +1,5 @@
-@testset "tasks" begin
+# Async tasks should consistently use the patch environment (if any) they started with
+@testset "async scope" begin
     c = Condition()
     ch = Channel{String}(1)
     f() = "original"
@@ -20,7 +21,12 @@
             @test (@mock f()) == "mocked"
 
             notify(c)
-            @test_broken take!(ch) == "original"
+            # https://github.com/JuliaLang/julia/pull/50958
+            if VERSION >= v"1.11.0-DEV.482"
+                @test take!(ch) == "original"
+            else
+                @test_broken take!(ch) == "original"
+            end
 
             # Task started inside patched context should call patched functions.
             @async background()
@@ -35,6 +41,11 @@
         end
 
         notify(c)
-        @test_broken take!(ch) == "mocked"
+        # https://github.com/JuliaLang/julia/pull/50958
+        if VERSION >= v"1.11.0-DEV.482"
+            @test take!(ch) == "mocked"
+        else
+            @test_broken take!(ch) == "mocked"
+        end
     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.apply(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)
+    # Patch environment has been reset back to original clean state
+    @test Mocking.PATCH_ENV[] == Mocking.PatchEnv()
 
     # Ensure that original behaviour is restored
     @test (@mock multiply(2)) == 3

test/runtests.jl

@@ -28,7 +28,7 @@ using Mocking: anon_morespecific, anonymous_signature, dispatch, type_morespecif
     include("args.jl")
     include("merge.jl")
     include("nested_apply.jl")
-    include("async.jl")
+    include("async-scope.jl")
     include("issues.jl")
     include("activate.jl")
     include("async-world-ages.jl")