feat(tests): port changes from neotest/nio

doc/nio.txt

@@ -694,25 +694,22 @@ Wrappers around plenary.nvim's test functions for writing async tests
    nio.sleep(10)
    assert.equals(10, notified)
  end)
+Fields~
+{it} `(fun(name: string, async_fun: fun()))`
+{before_each} `(fun(async_fun: fun()))`
+{after_each} `(fun(async_fun: fun()))`
 
-                                                                *nio.tests.it()*
-`it`({name}, {async_func})
-
-Parameters~
-{name} `(string)`
-{async_func} `(function)`
-
-                                                       *nio.tests.before_each()*
-`before_each`({async_func})
-
-Parameters~
-{async_func} `(function)`
-
-                                                        *nio.tests.after_each()*
-`after_each`({async_func})
+                                                *nio.tests.with_async_context()*
+`with_async_context`({async_func}, {...})
 
+the given function, applied to the remaining arguments, in an
+context.  The return value (or values) is the return value of
+asynchronous function.
 Parameters~
-{async_func} `(function)`
+{async_func} `(function)`  Function to execute
+{...} `(any)`  Arguments to `async_func`
+Return~
+`(any)` ...  Return values of `async_func`
 
 
  vim:tw=78:ts=8:noet:ft=help:norl:

lua/nio/tests.lua

@@ -21,15 +21,20 @@ local nio = {}
 ---    assert.equals(10, notified)
 ---  end)
 ---@class nio.tests
+---@field it fun(name: string, async_fun: fun())
+---@field before_each fun(async_fun: fun())
+---@field after_each fun(async_fun: fun())
 nio.tests = {}
 
 local with_timeout = function(func, timeout)
+  local success, err, results
   return function()
-    local success, err
-    local task = tasks.run(func, function(success_, err_)
+    local task = tasks.run(func, function(success_, ...)
       success = success_
       if not success_ then
-        err = err_
+        err = ...
+      else
+        results = { ... }
       end
     end)
 
@@ -42,23 +47,47 @@ local with_timeout = function(func, timeout)
     elseif not success then
       error(string.format("Test task failed with message:\n%s", err))
     end
+    return unpack(results)
   end
 end
 
----@param name string
----@param async_func function
-nio.tests.it = function(name, async_func)
-  it(name, with_timeout(async_func, tonumber(vim.env.PLENARY_TEST_TIMEOUT)))
-end
+local mt = {
+  __index = function(_table, key)
+    -- Hook functions from busted are only available within scope when the
+    -- test is defined, not when it is run, so we need to capture them
+    -- dynamically here.
+    local hook = getfenv(2)[key]
+    if not hook then
+      return nil
+    end
 
----@param async_func function
-nio.tests.before_each = function(async_func)
-  before_each(with_timeout(async_func))
-end
+    if key == "it" then
+      return function(name, async_func)
+        hook(name, with_timeout(async_func, tonumber(vim.env.PLENARY_TEST_TIMEOUT)))
+      end
+    elseif key == "before_each" or key == "after_each" then
+      return function(async_func)
+        hook(with_timeout(async_func))
+      end
+    end
+  end,
+}
 
----@param async_func function
-nio.tests.after_each = function(async_func)
-  after_each(with_timeout(async_func))
+setmetatable(nio.tests, mt)
+
+---Run the given function, applied to the remaining arguments, in an
+---asynchronous context.  The return value (or values) is the return value of
+---the asynchronous function.
+---@param async_func function  Function to execute
+---@param ... any  Arguments to `async_func`
+---@return any ...  Return values of `async_func`
+nio.tests.with_async_context = function(async_func, ...)
+  local args = { ... }
+  local thunk = function()
+    return async_func(unpack(args))
+  end
+  local result = with_timeout(thunk, tonumber(vim.env.PLENARY_TEST_TIMEOUT))()
+  return result
 end
 
 return nio.tests