Support custom records

Project.toml

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

src/ParallelTestRunner.jl

@@ -1,6 +1,7 @@
 module ParallelTestRunner
 
 export runtests, addworkers, addworker
+# public extract_flag!
 
 using Distributed
 using Dates
@@ -100,7 +101,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 +116,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 +130,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 +144,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 +159,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 +168,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 +194,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,9 +213,9 @@ end
 
 #
 # entry point
-#
+# 
 
-function runtest(::Type{TestRecord}, f, name, init_code, color)
+function runtest(::Type{TestRecord}, f, name, init_code, color, custom_args)
     function inner()
         # generate a temporary module to execute the tests in
         mod = @eval(Main, module $(gensym(name)) end)
@@ -466,7 +467,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 +791,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,99 @@ 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.runtest(::Type{CustomTestRecord}, f, name, init_code, color, (; say_hello))
+            function inner()
+                # 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)
+
+                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
+
+            @static if VERSION >= v"1.13.0-DEV.1044"
+                @with Test.TESTSET_PRINT_ENABLE => false begin
+                    inner()
+                end
+            else
+                old_print_setting = Test.TESTSET_PRINT_ENABLE[]
+                Test.TESTSET_PRINT_ENABLE[] = false
+                try
+                    inner()
+                finally
+                    Test.TESTSET_PRINT_ENABLE[] = old_print_setting
+                end
+            end
+        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