Support custom records

Project.toml

@@ -1,7 +1,7 @@
 name = "ParallelTestRunner"
 uuid = "d3525ed8-44d0-4b2c-a655-542cee43accc"
 authors = ["Valentin Churavy <v.churavy@gmail.com>"]
-version = "1.0.2"
+version = "1.1.0"
 
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"

src/ParallelTestRunner.jl

@@ -1,6 +1,9 @@
 module ParallelTestRunner
 
 export runtests, addworkers, addworker
+if VERSION >= v"1.11.0-DEV.469"
+    eval(Meta.parse("public extract_flag!"))
+end
 
 using Distributed
 using Dates
@@ -100,7 +103,7 @@ struct TestIOContext
     rss_align::Int
 end
 
-function test_IOContext(::Type{TestRecord}, stdout::IO, stderr::IO, lock::ReentrantLock, name_align::Int)
+function test_IOContext(::Type{<:AbstractTestRecord}, stdout::IO, stderr::IO, lock::ReentrantLock, name_align::Int)
     elapsed_align = textwidth("Time (s)")
     gc_align = textwidth("GC (s)")
     percent_align = textwidth("GC %")
@@ -115,7 +118,7 @@ function test_IOContext(::Type{TestRecord}, stdout::IO, stderr::IO, lock::Reentr
     )
 end
 
-function print_header(::Type{TestRecord}, ctx::TestIOContext, testgroupheader, workerheader)
+function print_header(::Type{<:AbstractTestRecord}, ctx::TestIOContext, testgroupheader, workerheader)
     lock(ctx.lock)
     try
         printstyled(ctx.stdout, " "^(ctx.name_align + textwidth(testgroupheader) - 3), " │ ")
@@ -129,7 +132,7 @@ function print_header(::Type{TestRecord}, ctx::TestIOContext, testgroupheader, w
     end
 end
 
-function print_test_started(::Type{TestRecord}, wrkr, test, ctx::TestIOContext)
+function print_test_started(::Type{<:AbstractTestRecord}, wrkr, test, ctx::TestIOContext)
     lock(ctx.lock)
     try
         printstyled(ctx.stdout, test, lpad("($wrkr)", ctx.name_align - textwidth(test) + 1, " "), " │", color = :white)
@@ -143,7 +146,7 @@ function print_test_started(::Type{TestRecord}, wrkr, test, ctx::TestIOContext)
     end
 end
 
-function print_test_finished(record::TestRecord, wrkr, test, ctx::TestIOContext)
+function print_test_finished(record::AbstractTestRecord, wrkr, test, ctx::TestIOContext)
     lock(ctx.lock)
     try
         printstyled(ctx.stdout, test, color = :white)
@@ -158,7 +161,7 @@ function print_test_finished(record::TestRecord, wrkr, test, ctx::TestIOContext)
         alloc_str = @sprintf("%5.2f", record.bytes / 2^20)
         printstyled(ctx.stdout, lpad(alloc_str, ctx.alloc_align, " "), " │ ", color = :white)
 
-        rss_str = @sprintf("%5.2f", record.rss / 2^20)
+        rss_str = @sprintf("%5.2f", memory_usage(record) / 2^20)
         printstyled(ctx.stdout, lpad(rss_str, ctx.rss_align, " "), " │\n", color = :white)
 
         flush(ctx.stdout)
@@ -167,7 +170,7 @@ function print_test_finished(record::TestRecord, wrkr, test, ctx::TestIOContext)
     end
 end
 
-function print_test_failed(record::TestRecord, wrkr, test, ctx::TestIOContext)
+function print_test_failed(record::AbstractTestRecord, wrkr, test, ctx::TestIOContext)
     lock(ctx.lock)
     try
         printstyled(ctx.stderr, test, color = :red)
@@ -193,7 +196,7 @@ function print_test_failed(record::TestRecord, wrkr, test, ctx::TestIOContext)
     end
 end
 
-function print_test_crashed(::Type{TestRecord}, wrkr, test, ctx::TestIOContext)
+function print_test_crashed(::Type{<:AbstractTestRecord}, wrkr, test, ctx::TestIOContext)
     lock(ctx.lock)
     try
         printstyled(ctx.stderr, test, color = :red)
@@ -212,58 +215,57 @@ end
 
 #
 # entry point
-#
-
-function runtest(::Type{TestRecord}, f, name, init_code, color)
-    function inner()
-        # generate a temporary module to execute the tests in
-        mod = @eval(Main, module $(gensym(name)) end)
-        @eval(mod, import ParallelTestRunner: Test, Random)
-        @eval(mod, using .Test, .Random)
-
-        Core.eval(mod, init_code)
+# 
 
-        data = @eval mod begin
-            GC.gc(true)
-            Random.seed!(1)
-
-            mktemp() do path, io
-                stats = redirect_stdio(stdout=io, stderr=io) do
-                    @timed try
-                        @testset $name begin
-                            $f
-                        end
-                    catch err
-                        isa(err, Test.TestSetException) || rethrow()
+function execute(::Type{TestRecord}, mod, f, name, color, custom_args)::TestRecord
+    data = @eval mod begin
+        GC.gc(true)
+        Random.seed!(1)
 
-                        # return the error to package it into a TestRecord
-                        err
+        mktemp() do path, io
+            stats = redirect_stdio(stdout=io, stderr=io) do
+                @timed try
+                    @testset $name begin
+                        $f
                     end
-                end
-                close(io)
-                output = read(path, String)
-                (; testset=stats.value, output, stats.time, stats.bytes, stats.gctime)
+                catch err
+                    isa(err, Test.TestSetException) || rethrow()
 
+                    # return the error to package it into a TestRecord
+                    err
+                end
             end
+            close(io)
+            output = read(path, String)
+            (; testset=stats.value, output, stats.time, stats.bytes, stats.gctime)
         end
+    end
 
-        # process results
-        rss = Sys.maxrss()
-        record = TestRecord(data..., rss)
+    # process results
+    rss = Sys.maxrss()
+    record = TestRecord(data..., rss)
 
-        GC.gc(true)
-        return record
-    end
+    GC.gc(true)
+    return record
+end
+
+function runtest(RecordType::Type{<:AbstractTestRecord}, f, name, init_code, color, custom_args)
+    # generate a temporary module to execute the tests in
+    mod = Core.eval(Main, Expr(:module, true, gensym(name), Expr(:block)))
+    @eval(mod, import ParallelTestRunner: Test, Random)
+    @eval(mod, using .Test, .Random)
+
+    Core.eval(mod, init_code)
 
     @static if VERSION >= v"1.13.0-DEV.1044"
         @with Test.TESTSET_PRINT_ENABLE => false begin
-            inner()
+            execute(RecordType, mod, f, name, color, custom_args)
         end
     else
         old_print_setting = Test.TESTSET_PRINT_ENABLE[]
         Test.TESTSET_PRINT_ENABLE[] = false
         try
-            inner()
+            execute(RecordType, mod, f, name, color, custom_args)
         finally
             Test.TESTSET_PRINT_ENABLE[] = old_print_setting
         end
@@ -466,7 +468,8 @@ Workers are automatically recycled when they exceed memory limits to prevent out
 issues during long test runs. The memory limit is set based on system architecture.
 """
 function runtests(mod::Module, ARGS; test_filter = Returns(true), RecordType = TestRecord,
-                  custom_tests::Dict{String, Expr}=Dict{String, Expr}(), init_code = :(),
+                  custom_tests::Dict{String, Expr}=Dict{String, Expr}(), init_code = :(), 
+                  custom_record_init = :(), custom_args = (;),
                   test_worker = Returns(nothing), stdout = Base.stdout, stderr = Base.stderr)
     #
     # set-up
@@ -789,8 +792,9 @@ function runtests(mod::Module, ARGS; test_filter = Returns(true), RecordType = T
                 put!(printer_channel, (:started, test, wrkr))
                 result = try
                     Distributed.remotecall_eval(Main, wrkr, :(import ParallelTestRunner))
+                    custom_record_init != :() && Distributed.remotecall_eval(Main, wrkr, custom_record_init)
                     remotecall_fetch(runtest, wrkr, RecordType, test_runners[test], test,
-                                              init_code, io_ctx.color)
+                                              init_code, io_ctx.color, custom_args)
                 catch ex
                     if isa(ex, InterruptException)
                         # the worker got interrupted, signal other tasks to stop

test/runtests.jl

@@ -75,6 +75,76 @@ end
     @test contains(str, "SUCCESS")
 end
 
+@testset "custom testrecord" begin
+    custom_record_init = quote
+        import ParallelTestRunner: Test
+        struct CustomTestRecord <: ParallelTestRunner.AbstractTestRecord
+            # TODO: Would it be better to wrap "ParallelTestRunner.TestRecord "
+            value::Any          # AbstractTestSet or TestSetException
+            output::String      # captured stdout/stderr
+
+            # stats
+            time::Float64
+            bytes::UInt64
+            gctime::Float64
+            rss::UInt64
+        end
+        function ParallelTestRunner.memory_usage(rec::CustomTestRecord)
+            return rec.rss
+        end
+        function ParallelTestRunner.test_IOContext(::Type{CustomTestRecord}, stdout::IO, stderr::IO, lock::ReentrantLock, name_align::Int64)
+            return ParallelTestRunner.test_IOContext(ParallelTestRunner.TestRecord, stdout, stderr, lock, name_align)
+        end
+        function ParallelTestRunner.execute(::Type{CustomTestRecord}, mod, f, name, color, (; say_hello))
+            data = @eval mod begin
+                GC.gc(true)
+                Random.seed!(1)
+
+                mktemp() do path, io
+                    stats = redirect_stdio(stdout=io, stderr=io) do
+                        @timed try
+                            # Since we are in a double quote we need to use this form to escape `$`
+                            if $(Expr(:$, :say_hello))
+                                println("Hello from test '" * $(Expr(:$, :name)) * "'")
+                            end
+                            @testset $(Expr(:$, :name)) begin
+                                $(Expr(:$, :f))
+                            end
+                        catch err
+                            isa(err, Test.TestSetException) || rethrow()
+
+                            # return the error to package it into a TestRecord
+                            err
+                        end
+                    end
+                    close(io)
+                    output = read(path, String)
+                    (; testset=stats.value, output, stats.time, stats.bytes, stats.gctime)
+
+                end
+            end
+
+            # process results
+            rss = Sys.maxrss()
+            record = CustomTestRecord(data..., rss)
+
+            GC.gc(true)
+            return record
+        end
+    end # quote
+    eval(custom_record_init)
+
+    io = IOBuffer()
+
+    runtests(ParallelTestRunner, ["--verbose"]; custom_record_init, RecordType=CustomTestRecord, custom_args=(; say_hello=true), stdout=io, stderr=io)
+    str = String(take!(io))
+
+    @test contains(str, r"basic .+ started at")
+    @test contains(str, r"Hello from test 'basic'")
+    @test contains(str, "SUCCESS")
+end
+
+
 @testset "failing test" begin
     custom_tests = Dict(
         "failing test" => quote