From d0cca64a7c88405947019026b7ba441698ac22ca Mon Sep 17 00:00:00 2001
From: Nathan Daly <nathan.daly@relational.ai>
Date: Tue, 28 Nov 2023 17:31:06 -0700
Subject: [PATCH] test_persistent_tasks: Add an optional `expr` to run in the
 precompile package

---
 src/persistent_tasks.jl | 36 +++++++++++++++++++++++++++---------
 1 file changed, 27 insertions(+), 9 deletions(-)

diff --git a/src/persistent_tasks.jl b/src/persistent_tasks.jl
index 85092d1b..c229e4bf 100644
--- a/src/persistent_tasks.jl
+++ b/src/persistent_tasks.jl
@@ -1,10 +1,14 @@
 """
-    Aqua.test_persistent_tasks(package)
+    Aqua.test_persistent_tasks(package, [expr])
 
 Test whether loading `package` creates persistent `Task`s
 which may block precompilation of dependent packages.
 See also [`Aqua.find_persistent_tasks_deps`](@ref).
 
+If you provide an optional `expr`, this tests whether loading `package` and running `expr`
+creates persistent `Task`s. For example, you might start and shutdown a web server, and
+this will test that there aren't any persistent `Task`s.
+
 # Motivation
 
 Julia 1.10 and higher wait for all running `Task`s to finish
@@ -37,6 +41,18 @@ fails to precompile: `using PkgA` runs `PkgA.__init__()`, which
 leaves the `Timer` `Task` running, and that causes precompilation
 of `PkgB` to hang.
 
+# Example invocations
+
+```julia
+Aqua.test_persistent_tasks(MyPackage)
+
+Aqua.test_persistent_tasks(MyPackage, quote
+    # Code to run after loading MyPackage
+    server = MyPackage.start_server()
+    MyPackage.stop_server!(server)
+end)
+```
+
 # How the test works
 
 This test works by launching a Julia process that tries to precompile a
@@ -66,6 +82,7 @@ On Julia version 1.9 and before, this test always succeeds.
 
 # Arguments
 - `package`: a top-level `Module` or `Base.PkgId`.
+- `expr = nothing`: An expression to run in the precompile package.
 
 # Keyword Arguments
 - `broken::Bool = false`: If true, it uses `@test_broken` instead of
@@ -74,22 +91,22 @@ On Julia version 1.9 and before, this test always succeeds.
   package before forcibly shutting down the precompilation process (triggering
   a test failure).
 """
-function test_persistent_tasks(package::PkgId; broken::Bool = false, kwargs...)
+function test_persistent_tasks(package::PkgId, expr=nothing; broken::Bool = false, kwargs...)
     if broken
-        @test_broken !has_persistent_tasks(package; kwargs...)
+        @test_broken !has_persistent_tasks(package, expr; kwargs...)
     else
-        @test !has_persistent_tasks(package; kwargs...)
+        @test !has_persistent_tasks(package, expr; kwargs...)
     end
 end
 
-function test_persistent_tasks(package::Module; kwargs...)
-    test_persistent_tasks(PkgId(package); kwargs...)
+function test_persistent_tasks(package::Module, expr=nothing; kwargs...)
+    test_persistent_tasks(PkgId(package), expr; kwargs...)
 end
 
-function has_persistent_tasks(package::PkgId; tmax = 10)
+function has_persistent_tasks(package::PkgId, expr=nothing; tmax = 10)
     root_project_path, found = root_project_toml(package)
     found || error("Unable to locate Project.toml")
-    return !precompile_wrapper(root_project_path, tmax)
+    return !precompile_wrapper(root_project_path, tmax, expr)
 end
 
 """
@@ -117,7 +134,7 @@ function find_persistent_tasks_deps(package::Module; kwargs...)
     find_persistent_tasks_deps(PkgId(package); kwargs...)
 end
 
-function precompile_wrapper(project, tmax)
+function precompile_wrapper(project, tmax, expr)
     if VERSION < v"1.10.0-"
         return true
     end
@@ -141,6 +158,7 @@ function precompile_wrapper(project, tmax)
                 """
 module $wrappername
 using $pkgname
+$expr
 # Signal Aqua from the precompilation process that we've finished loading the package
 open("$(escape_string(statusfile))", "w") do io
     println(io, "done")