From d86658f9d9f4ca1bd700dd35418504bcf9a5fa3b Mon Sep 17 00:00:00 2001
From: zeertzjq <zeertzjq@outlook.com>
Date: Wed, 17 Sep 2025 08:21:11 +0800
Subject: [PATCH] test(lsp): make async format test work properly (#35794)

Overriding vim.lsp.handlers['textDocument/formatting'] doesn't work here
because fake_lsp_server_setup() uses a table with __index to specify
client handlers, which takes priority over vim.lsp.handlers[], and as a
result the overridden handler is never called, and the test ends before
the vim.wait() even finishes.

Instead, set a global variable from the handler that is actually reached
(by vim.rpcrequest() from client handler), and avoid stopping the event
loop too early.

(cherry picked from commit 8d5476691cd177ee2166bc8f4c64679292cc7559)
---
 test/functional/plugin/lsp/testutil.lua | 33 ++++++++++++++-----------
 test/functional/plugin/lsp_spec.lua     | 14 ++++-------
 2 files changed, 23 insertions(+), 24 deletions(-)

diff --git a/test/functional/plugin/lsp/testutil.lua b/test/functional/plugin/lsp/testutil.lua
index 5edcc00845..35ee7741ca 100644
--- a/test/functional/plugin/lsp/testutil.lua
+++ b/test/functional/plugin/lsp/testutil.lua
@@ -199,31 +199,34 @@ function M.test_rpc_server(config)
   })
   --- @type integer, integer
   local code, signal
+  local busy = 0
+  local exited = false
   local function on_request(method, args)
-    if method == 'setup' then
-      if config.on_setup then
-        config.on_setup()
-      end
-      return NIL
+    busy = busy + 1
+    if method == 'setup' and config.on_setup then
+      config.on_setup()
     end
-    if method == 'init' then
-      if config.on_init then
-        config.on_init(client, unpack(args))
-      end
-      return NIL
+    if method == 'init' and config.on_init then
+      config.on_init(client, unpack(args))
     end
-    if method == 'handler' then
-      if config.on_handler then
-        config.on_handler(unpack(args))
-      end
+    if method == 'handler' and config.on_handler then
+      config.on_handler(unpack(args))
+    end
+    busy = busy - 1
+    if busy == 0 and exited then
+      stop()
     end
     return NIL
   end
   local function on_notify(method, args)
     if method == 'exit' then
       code, signal = unpack(args)
-      return stop()
+      exited = true
+      if busy == 0 then
+        stop()
+      end
     end
+    return NIL
   end
   --  TODO specify timeout?
   --  run(on_request, on_notify, nil, 1000)
diff --git a/test/functional/plugin/lsp_spec.lua b/test/functional/plugin/lsp_spec.lua
index 638918c629..28f159a3f9 100644
--- a/test/functional/plugin/lsp_spec.lua
+++ b/test/functional/plugin/lsp_spec.lua
@@ -5132,22 +5132,18 @@ describe('LSP', function()
                 notify_msg = msg
               end
 
-              local handler = vim.lsp.handlers['textDocument/formatting']
-              local handler_called = false
-              vim.lsp.handlers['textDocument/formatting'] = function()
-                handler_called = true
-              end
-
+              _G.handler_called = false
               vim.lsp.buf.format({ bufnr = bufnr, async = true })
               vim.wait(1000, function()
-                return handler_called
+                return _G.handler_called
               end)
 
               vim.notify = notify
-              vim.lsp.handlers['textDocument/formatting'] = handler
-              return { notify = notify_msg, handler_called = handler_called }
+              return { notify_msg = notify_msg, handler_called = _G.handler_called }
             end)
             eq({ handler_called = true }, result)
+          elseif ctx.method == 'textDocument/formatting' then
+            exec_lua('_G.handler_called = true')
           elseif ctx.method == 'shutdown' then
             client:stop()
           end