use ScopedValues for TestSets

Author: simonbyrne

stdlib/Test/src/Test.jl

@@ -1004,7 +1004,6 @@ end
 A simple fallback test set that throws immediately on a failure.
 """
 struct FallbackTestSet <: AbstractTestSet end
-fallback_testset = FallbackTestSet()
 
 struct FallbackTestSetException <: Exception
     msg::String
@@ -1639,22 +1638,11 @@ function testset_context(args, ex, source)
     else
         error("Malformed `let` expression is given")
     end
-    reverse!(contexts)
-
-    test_ex = ex.args[2]
-
-    ex.args[2] = quote
-        $(map(contexts) do context
-            :($push_testset($(ContextTestSet)($(QuoteNode(context)), $context; $options...)))
-        end...)
-        try
-            $(test_ex)
-        finally
-            $(map(_->:($pop_testset()), contexts)...)
-        end
+    test_ex = esc(ex.args[2])
+    for context in reverse(contexts)
+        test_ex = :(@with_testset(ContextTestSet($(QuoteNode(context)), $(esc(context)); $(esc(options))...), $test_ex))
     end
-
-    return esc(ex)
+    return test_ex
 end
 
 """
@@ -1687,35 +1675,34 @@ function testset_beginend_call(args, tests, source)
         else
             $(testsettype)($desc; $options...)
         end
-        push_testset(ts)
-        # we reproduce the logic of guardseed, but this function
-        # cannot be used as it changes slightly the semantic of @testset,
-        # by wrapping the body in a function
-        local default_rng_orig = copy(default_rng())
-        local tls_seed_orig = copy(Random.get_tls_seed())
-        try
-            # default RNG is reset to its state from last `seed!()` to ease reproduce a failed test
-            copy!(Random.default_rng(), tls_seed_orig)
-            let
-                $(esc(tests))
-            end
-        catch err
-            err isa InterruptException && rethrow()
-            # something in the test block threw an error. Count that as an
-            # error in this test set
-            trigger_test_failure_break(err)
-            if err isa FailFastError
-                get_testset_depth() > 1 ? rethrow() : failfast_print()
-            else
-                record(ts, Error(:nontest_error, Expr(:tuple), err, Base.current_exceptions(), $(QuoteNode(source))))
+        @with_testset ts begin
+            # we reproduce the logic of guardseed, but this function
+            # cannot be used as it changes slightly the semantic of @testset,
+            # by wrapping the body in a function
+            local default_rng_orig = copy(default_rng())
+            local tls_seed_orig = copy(Random.get_tls_seed())
+            try
+                # default RNG is reset to its state from last `seed!()` to ease reproduce a failed test
+                copy!(Random.default_rng(), tls_seed_orig)
+                let
+                    $(esc(tests))
+                end
+            catch err
+                err isa InterruptException && rethrow()
+                # something in the test block threw an error. Count that as an
+                # error in this test set
+                trigger_test_failure_break(err)
+                if err isa FailFastError
+                    get_testset_depth() > 1 ? rethrow() : failfast_print()
+                else
+                    record(ts, Error(:nontest_error, Expr(:tuple), err, Base.current_exceptions(), $(QuoteNode(source))))
+                end
+            finally
+                copy!(default_rng(), default_rng_orig)
+                copy!(Random.get_tls_seed(), tls_seed_orig)
             end
-        finally
-            copy!(default_rng(), default_rng_orig)
-            copy!(Random.get_tls_seed(), tls_seed_orig)
-            pop_testset()
-            ret = finish(ts)
         end
-        ret
+        ts
     end
     # preserve outer location if possible
     if tests isa Expr && tests.head === :block && !isempty(tests.args) && tests.args[1] isa LineNumberNode
@@ -1771,52 +1758,38 @@ function testset_forloop(args, testloop, source)
         _check_testset($testsettype, $(QuoteNode(testsettype.args[1])))
         # Trick to handle `break` and `continue` in the test code before
         # they can be handled properly by `finally` lowering.
-        if !first_iteration
-            pop_testset()
-            finish_errored = true
-            push!(arr, finish(ts))
-            finish_errored = false
-            copy!(default_rng(), tls_seed_orig)
-        end
         ts = if ($testsettype === $DefaultTestSet) && $(isa(source, LineNumberNode))
             $(testsettype)($desc; source=$(QuoteNode(source.file)), $options...)
         else
             $(testsettype)($desc; $options...)
         end
-        push_testset(ts)
-        first_iteration = false
-        try
-            $(esc(tests))
-        catch err
-            err isa InterruptException && rethrow()
-            # Something in the test block threw an error. Count that as an
-            # error in this test set
-            trigger_test_failure_break(err)
-            if !isa(err, FailFastError)
-                record(ts, Error(:nontest_error, Expr(:tuple), err, Base.current_exceptions(), $(QuoteNode(source))))
+        @with_testset ts begin
+            try
+                # default RNG is reset to its state from last `seed!()` to ease reproduce a failed test
+                copy!(Random.default_rng(), tls_seed_orig)
+                $(esc(tests))
+            catch err
+                err isa InterruptException && rethrow()
+                # Something in the test block threw an error. Count that as an
+                # error in this test set
+                trigger_test_failure_break(err)
+                if !isa(err, FailFastError)
+                    record(ts, Error(:nontest_error, Expr(:tuple), err, Base.current_exceptions(), $(QuoteNode(source))))
+                end
+            finally
+                copy!(default_rng(), default_rng_orig)
+                copy!(Random.get_tls_seed(), tls_seed_orig)
+                push!(arr, ts)
             end
         end
     end
     quote
         local arr = Vector{Any}()
-        local first_iteration = true
-        local ts
-        local finish_errored = false
         local default_rng_orig = copy(default_rng())
         local tls_seed_orig = copy(Random.get_tls_seed())
-        copy!(Random.default_rng(), tls_seed_orig)
-        try
-            let
-                $(Expr(:for, Expr(:block, [esc(v) for v in loopvars]...), blk))
-            end
-        finally
-            # Handle `return` in test body
-            if !first_iteration && !finish_errored
-                pop_testset()
-                push!(arr, finish(ts))
-            end
-            copy!(default_rng(), default_rng_orig)
-            copy!(Random.get_tls_seed(), tls_seed_orig)
+        local ts
+        let
+            $(Expr(:for, Expr(:block, [esc(v) for v in loopvars]...), blk))
         end
         arr
     end
@@ -1855,39 +1828,27 @@ end
 #-----------------------------------------------------------------------
 # Various helper methods for test sets
 
+const CURRENT_TESTSET = ScopedValue{AbstractTestSet}(FallbackTestSet())
+const TESTSET_DEPTH = ScopedValue{Int}(0)
+
+macro with_testset(ts, expr)
+    quote
+        ts = $(esc(ts))
+        Expr(:tryfinally,
+            @with(CURRENT_TESTSET => ts, TESTSET_DEPTH => get_testset_depth() + 1, $(esc(expr))),
+            finish(ts)
+        )
+    end
+end
+
 """
     get_testset()
 
 Retrieve the active test set from the task's local storage. If no
 test set is active, use the fallback default test set.
 """
 function get_testset()
-    testsets = get(task_local_storage(), :__BASETESTNEXT__, AbstractTestSet[])
-    return isempty(testsets) ? fallback_testset : testsets[end]
-end
-
-"""
-    push_testset(ts::AbstractTestSet)
-
-Adds the test set to the `task_local_storage`.
-"""
-function push_testset(ts::AbstractTestSet)
-    testsets = get(task_local_storage(), :__BASETESTNEXT__, AbstractTestSet[])
-    push!(testsets, ts)
-    setindex!(task_local_storage(), testsets, :__BASETESTNEXT__)
-end
-
-"""
-    pop_testset()
-
-Pops the last test set added to the `task_local_storage`. If there are no
-active test sets, returns the fallback default test set.
-"""
-function pop_testset()
-    testsets = get(task_local_storage(), :__BASETESTNEXT__, AbstractTestSet[])
-    ret = isempty(testsets) ? fallback_testset : pop!(testsets)
-    setindex!(task_local_storage(), testsets, :__BASETESTNEXT__)
-    return ret
+    something(Base.ScopedValues.get(CURRENT_TESTSET))
 end
 
 """
@@ -1896,10 +1857,10 @@ end
 Return the number of active test sets, not including the default test set
 """
 function get_testset_depth()
-    testsets = get(task_local_storage(), :__BASETESTNEXT__, AbstractTestSet[])
-    return length(testsets)
+    something(Base.ScopedValues.get(TESTSET_DEPTH))
 end
 
+
 _args_and_call(args...; kwargs...) = (args[1:end-1], kwargs, args[end](args[1:end-1]...; kwargs...))
 _materialize_broadcasted(f, args...) = Broadcast.materialize(Broadcast.broadcasted(f, args...))
 

test/runtests.jl

@@ -379,55 +379,55 @@ cd(@__DIR__) do
     Test.TESTSET_PRINT_ENABLE[] = false
     o_ts = Test.DefaultTestSet("Overall")
     o_ts.time_end = o_ts.time_start + o_ts_duration # manually populate the timing
-    Test.push_testset(o_ts)
-    completed_tests = Set{String}()
-    for (testname, (resp,), duration) in results
-        push!(completed_tests, testname)
-        if isa(resp, Test.DefaultTestSet)
-            resp.time_end = resp.time_start + duration
-            Test.push_testset(resp)
-            Test.record(o_ts, resp)
-            Test.pop_testset()
-        elseif isa(resp, Test.TestSetException)
-            fake = Test.DefaultTestSet(testname)
-            fake.time_end = fake.time_start + duration
-            for i in 1:resp.pass
-                Test.record(fake, Test.Pass(:test, nothing, nothing, nothing, LineNumberNode(@__LINE__, @__FILE__)))
-            end
-            for i in 1:resp.broken
-                Test.record(fake, Test.Broken(:test, nothing))
-            end
-            for t in resp.errors_and_fails
-                Test.record(fake, t)
+    Test.@with_testset o_ts begin
+        completed_tests = Set{String}()
+        for (testname, (resp,), duration) in results
+            push!(completed_tests, testname)
+            if isa(resp, Test.DefaultTestSet)
+                resp.time_end = resp.time_start + duration
+                Test.@with_testset resp begin
+                    Test.record(o_ts, resp)
+                end
+            elseif isa(resp, Test.TestSetException)
+                fake = Test.DefaultTestSet(testname)
+                fake.time_end = fake.time_start + duration
+                for i in 1:resp.pass
+                    Test.record(fake, Test.Pass(:test, nothing, nothing, nothing, LineNumberNode(@__LINE__, @__FILE__)))
+                end
+                for i in 1:resp.broken
+                    Test.record(fake, Test.Broken(:test, nothing))
+                end
+                for t in resp.errors_and_fails
+                    Test.record(fake, t)
+                end
+                Test.@with_testset fake begin
+                    Test.record(o_ts, fake)
+                end
+            else
+                if !isa(resp, Exception)
+                    resp = ErrorException(string("Unknown result type : ", typeof(resp)))
+                end
+                # If this test raised an exception that is not a remote testset exception,
+                # i.e. not a RemoteException capturing a TestSetException that means
+                # the test runner itself had some problem, so we may have hit a segfault,
+                # deserialization errors or something similar.  Record this testset as Errored.
+                fake = Test.DefaultTestSet(testname)
+                fake.time_end = fake.time_start + duration
+                Test.record(fake, Test.Error(:nontest_error, testname, nothing, Any[(resp, [])], LineNumberNode(1)))
+                Test.@with_testset fake begin
+                    Test.record(o_ts, fake)
+                end
             end
-            Test.push_testset(fake)
-            Test.record(o_ts, fake)
-            Test.pop_testset()
-        else
-            if !isa(resp, Exception)
-                resp = ErrorException(string("Unknown result type : ", typeof(resp)))
+        end
+        for test in all_tests
+            (test in completed_tests) && continue
+            fake = Test.DefaultTestSet(test)
+            Test.record(fake, Test.Error(:test_interrupted, test, nothing, [("skipped", [])], LineNumberNode(1)))
+            Test.@with_testset fake begin
+                Test.record(o_ts, fake)
             end
-            # If this test raised an exception that is not a remote testset exception,
-            # i.e. not a RemoteException capturing a TestSetException that means
-            # the test runner itself had some problem, so we may have hit a segfault,
-            # deserialization errors or something similar.  Record this testset as Errored.
-            fake = Test.DefaultTestSet(testname)
-            fake.time_end = fake.time_start + duration
-            Test.record(fake, Test.Error(:nontest_error, testname, nothing, Any[(resp, [])], LineNumberNode(1)))
-            Test.push_testset(fake)
-            Test.record(o_ts, fake)
-            Test.pop_testset()
         end
     end
-    for test in all_tests
-        (test in completed_tests) && continue
-        fake = Test.DefaultTestSet(test)
-        Test.record(fake, Test.Error(:test_interrupted, test, nothing, [("skipped", [])], LineNumberNode(1)))
-        Test.push_testset(fake)
-        Test.record(o_ts, fake)
-        Test.pop_testset()
-    end
-
     if Base.get_bool_env("CI", false)
         @info "Writing test result data to $(@__DIR__)"
         write_testset_json_files(@__DIR__, o_ts)