1 files changed, 16 insertions, 4 deletions
diff --git a/test/functional/helpers.lua b/test/functional/helpers.lua
index 0c616e73fb..8c5a60657a 100644
--- a/test/functional/helpers.lua
+++ b/test/functional/helpers.lua
@@ -271,10 +271,22 @@ function module.command(cmd)
 end
 
 
--- use for commands which expect nvim to quit
-function module.expect_exit(...)
-  eq("EOF was received from Nvim. Likely the Nvim process crashed.",
-     module.pcall_err(...))
+-- Use for commands which expect nvim to quit.
+-- The first argument can also be a timeout.
+function module.expect_exit(fn_or_timeout, ...)
+  local eof_err_msg = 'EOF was received from Nvim. Likely the Nvim process crashed.'
+  if type(fn_or_timeout) == 'function' then
+    eq(eof_err_msg, module.pcall_err(fn_or_timeout, ...))
+  else
+    eq(eof_err_msg, module.pcall_err(function(timeout, fn, ...)
+      fn(...)
+      while session:next_message(timeout) do
+      end
+      if session.eof_err then
+        error(session.eof_err[2])
+      end
+    end, fn_or_timeout, ...))
+  end
 end
 
 -- Evaluates a VimL expression.