init_worker_code with custom test_worker

Author: christiangnrd

docs/src/advanced.md

@@ -166,6 +166,11 @@ The `test_worker` function receives the test name and should return either:
 - A worker object (from [`addworker`](@ref)) for tests that need special configuration
 - `nothing` to use the default worker pool
 
+!!! note
+    If your test suite uses both a `test_worker` function and `init_worker_code` as described in a prior section,
+    `test_worker` must also take in `init_worker_code` as a second argument. You are responsible for passing it to
+    [`addworker`](@ref) if your `init_code` depends on any `init_worker_code` definitions.
+
 ## Custom Arguments
 
 If your package needs to accept its own command-line arguments in addition to `ParallelTestRunner`'s options, use [`parse_args`](@ref) with custom flags:

src/ParallelTestRunner.jl

@@ -686,7 +686,7 @@ Several keyword arguments are also supported:
 - `init_code`: Code use to initialize each test's sandbox module (e.g., import auxiliary
   packages, define constants, etc).
 - `init_worker_code`: Code use to initialize each worker. This is run only once per worker instead of once per test.
-- `test_worker`: Optional function that takes a test name and returns a specific worker.
+- `test_worker`: Optional function that takes a test name and `init_worker_code` if `init_worker_code` is defined and returns a specific worker.
   When returning `nothing`, the test will be assigned to any available default worker.
 - `stdout` and `stderr`: I/O streams to write to (default: `Base.stdout` and `Base.stderr`)
 
@@ -996,11 +996,16 @@ function runtests(mod::Module, args::ParsedArgs;
                     test, test_t0
                 end
 
-                # if a worker failed, spawn a new one
-                wrkr = test_worker(test)
-                if wrkr ===  nothing
+                # pass in init_worker_code to custom worker function if defined
+                wrkr = if init_worker_code == :()
+                    test_worker(test)
+                else
+                    test_worker(test, init_worker_code)
+                end
+                if wrkr === nothing
                     wrkr = p
                 end
+                # if a worker failed, spawn a new one
                 if wrkr === nothing || !Malt.isrunning(wrkr)
                     wrkr = p = addworker(; init_worker_code, io_ctx.color)
                 end

test/runtests.jl

@@ -134,6 +134,52 @@ end
     @test contains(str, "SUCCESS")
 end
 
+@testset "custom worker with `init_worker_code`" begin
+    init_worker_code = quote
+        should_be_defined() = true
+    end
+    init_code = quote
+        using Test
+        import ..should_be_defined
+    end
+    function test_worker(name, init_worker_code)
+        if name == "needs env var"
+            return addworker(env = ["SPECIAL_ENV_VAR" => "42"]; init_worker_code)
+        elseif name == "threads/2"
+            return addworker(exeflags = ["--threads=2"]; init_worker_code)
+        end
+        return nothing
+    end
+    testsuite = Dict(
+        "needs env var" => quote
+            @test ENV["SPECIAL_ENV_VAR"] == "42"
+            @test should_be_defined()
+        end,
+        "doesn't need env var" => quote
+            @test !haskey(ENV, "SPECIAL_ENV_VAR")
+            @test should_be_defined()
+        end,
+        "threads/1" => quote
+            @test Base.Threads.nthreads() == 1
+            @test should_be_defined()
+        end,
+        "threads/2" => quote
+            @test Base.Threads.nthreads() == 2
+            @test should_be_defined()
+        end
+    )
+
+    io = IOBuffer()
+    runtests(ParallelTestRunner, ["--verbose"]; test_worker, init_code, init_worker_code, testsuite, stdout=io, stderr=io)
+
+    str = String(take!(io))
+    @test contains(str, r"needs env var .+ started at")
+    @test contains(str, r"doesn't need env var .+ started at")
+    @test contains(str, r"threads/1 .+ started at")
+    @test contains(str, r"threads/2 .+ started at")
+    @test contains(str, "SUCCESS")
+end
+
 @testset "failing test" begin
     testsuite = Dict(
         "failing test" => quote