From 0dccb5ebcf00de245e2060fa17822a9464f5c41a Mon Sep 17 00:00:00 2001
From: rcarriga <rcarriga@tcd.ie>
Date: Sat, 28 Dec 2024 08:39:04 +0000
Subject: [PATCH] fix(strategies): handle quick exit

---
 .../client/strategies/integrated/init.lua     | 19 +++++++++++++++++--
 lua/neotest/consumers/output_panel/init.lua   |  8 +++++++-
 lua/neotest/consumers/status.lua              |  2 +-
 lua/neotest/consumers/watch/init.lua          |  8 +++++---
 tests/unit/consumers/output_panel_spec.lua    |  6 +++---
 5 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/lua/neotest/client/strategies/integrated/init.lua b/lua/neotest/client/strategies/integrated/init.lua
index b4d37fb..1a64c8d 100644
--- a/lua/neotest/client/strategies/integrated/init.lua
+++ b/lua/neotest/client/strategies/integrated/init.lua
@@ -13,6 +13,7 @@ return function(spec)
   local env, cwd = spec.env, spec.cwd
 
   local finish_future = nio.control.future()
+  local output_finish_future = nio.control.future()
   local result_code = nil
   local command = spec.command
   local output_accum = FanoutAccum(function(prev, new)
@@ -40,6 +41,10 @@ return function(spec)
     height = spec.strategy.height,
     width = spec.strategy.width,
     on_stdout = function(_, data)
+      if #data == 1 and data[1] == "" then
+        output_finish_future.set()
+        return
+      end
       output_accum:push(table.concat(data, "\n"))
     end,
     on_exit = function(_, code)
@@ -69,7 +74,7 @@ return function(spec)
         queue.put_nowait(d)
       end)
       return function()
-        local data = nio.first({ queue.get, finish_future.wait })
+        local data = nio.first({ queue.get, output_finish_future.wait })
         if data then
           return data
         end
@@ -114,7 +119,17 @@ return function(spec)
     end,
     result = function()
       if result_code == nil then
-        finish_future:wait()
+        finish_future.wait()
+        if not output_finish_future.is_set() then
+          -- jobstart doesn't necessarily call on_stdout if the process
+          -- stops quickly, so add a timeout to prevent deadlock
+          nio.first({
+            output_finish_future.wait,
+            function()
+              nio.sleep(100)
+            end,
+          })
+        end
       end
       local close_err = nio.uv.fs_close(output_fd)
       assert(not close_err, close_err)
diff --git a/lua/neotest/consumers/output_panel/init.lua b/lua/neotest/consumers/output_panel/init.lua
index 32b05e6..9a2b8ad 100644
--- a/lua/neotest/consumers/output_panel/init.lua
+++ b/lua/neotest/consumers/output_panel/init.lua
@@ -55,7 +55,13 @@ local init = function(client)
     for file, _ in pairs(files_to_read) do
       local output = lib.files.read(file)
       local dos_newlines = string.find(output, "\r\n") ~= nil
-      if not pcall(nio.api.nvim_chan_send, chan, dos_newlines and output or output:gsub("\n", "\r\n")) then
+      if
+        not pcall(
+          nio.api.nvim_chan_send,
+          chan,
+          dos_newlines and output or output:gsub("\n", "\r\n")
+        )
+      then
         lib.notify(("Error sending output to term channel: %s"):format(chan), vim.log.levels.ERROR)
         chan = nil
         break
diff --git a/lua/neotest/consumers/status.lua b/lua/neotest/consumers/status.lua
index ff4e45e..cdfe403 100644
--- a/lua/neotest/consumers/status.lua
+++ b/lua/neotest/consumers/status.lua
@@ -45,7 +45,7 @@ local function init(client)
           virt_text = {
             { statuses[status].text .. " ", statuses[status].texthl },
           },
-          hl_mode = "combine"
+          hl_mode = "combine",
         })
       end
     end
diff --git a/lua/neotest/consumers/watch/init.lua b/lua/neotest/consumers/watch/init.lua
index 69ddd66..1bdd6a9 100644
--- a/lua/neotest/consumers/watch/init.lua
+++ b/lua/neotest/consumers/watch/init.lua
@@ -32,9 +32,11 @@ neotest.watch = {}
 local init = function(client)
   client.listeners.discover_positions = function(_, tree)
     for _, watcher in pairs(watchers) do
-      if watcher.tree:data().path == tree:data().path
-        and not watcher.discover_positions_event.is_set() then
-          watcher.discover_positions_event.set()
+      if
+        watcher.tree:data().path == tree:data().path
+        and not watcher.discover_positions_event.is_set()
+      then
+        watcher.discover_positions_event.set()
       end
     end
   end
diff --git a/tests/unit/consumers/output_panel_spec.lua b/tests/unit/consumers/output_panel_spec.lua
index 9f3351b..0d1f40d 100644
--- a/tests/unit/consumers/output_panel_spec.lua
+++ b/tests/unit/consumers/output_panel_spec.lua
@@ -23,7 +23,7 @@ describe("neotest consumer - output_panel", function()
   local files
   local dirs = { dir }
   local notify
-  local notify_msg = ''
+  local notify_msg = ""
 
   ---@return neotest.Tree
   local get_pos = function(...)
@@ -136,7 +136,7 @@ describe("neotest consumer - output_panel", function()
     for _, buf in ipairs(vim.api.nvim_list_bufs()) do
       vim.api.nvim_buf_delete(buf, { force = true })
     end
-    notify_msg = ''
+    notify_msg = ""
   end)
 
   describe("user forcefully closes the panel", function()
@@ -165,7 +165,7 @@ describe("neotest consumer - output_panel", function()
         assert.has_no_error(function()
           client:run_tree(tree, { strategy = mock_strategy })
         end)
-        assert.is_not.matches('Error sending output to term channel:', notify_msg)
+        assert.is_not.matches("Error sending output to term channel:", notify_msg)
       end)
       exit_future_2.set()
     end)