neovim/test/functional/plugin/lsp/testutil.lua

local n = require('test.functional.testnvim')()

local clear = n.clear
local exec_lua = n.exec_lua
local run = n.run
local stop = n.stop
local api = n.api
local NIL = vim.NIL

local M = {}

function M.clear_notrace()
  -- problem: here be dragons
  -- solution: don't look too closely for dragons
  clear {
    env = {
      NVIM_LUA_NOTRACK = '1',
      NVIM_APPNAME = 'nvim_lsp_test',
      VIMRUNTIME = os.getenv 'VIMRUNTIME',
    },
  }
end

M.create_tcp_echo_server = function()
  --- Create a TCP server that echos the first message it receives.
  --- @param host string
  --- @return integer
  function _G._create_tcp_server(host)
    local uv = vim.uv
    local server = assert(uv.new_tcp())
    local on_read = require('vim.lsp.rpc').create_read_loop(
      function(body)
        vim.rpcnotify(1, 'body', body)
      end,
      nil,
      function(err, code)
        vim.rpcnotify(1, 'error', err, code)
      end
    )
    server:bind(host, 0)
    server:listen(127, function(e)
      assert(not e, e)
      local socket = assert(uv.new_tcp())
      server:accept(socket)
      socket:read_start(function(err, chunk)
        on_read(err, chunk)
        socket:shutdown()
        socket:close()
        server:shutdown()
        server:close()
      end)
    end)
    return server:getsockname().port
  end
  function _G._send_msg_to_server(msg)
    local port = _G._create_tcp_server('127.0.0.1')
    local client = assert(vim.uv.new_tcp())
    client:connect('127.0.0.1', port, function()
      client:write(msg, function()
        client:shutdown()
        client:close()
      end)
    end)
  end
end

M.create_server_definition = function()
  ---@class test.functional.plugin.lsp.testutil._create_server.Opts
  ---@field capabilities lsp.ServerCapabilities?
  ---@field handlers table<vim.lsp.protocol.Method.ClientToServer.Request, fun(method: vim.lsp.protocol.Method.ClientToServer.Request, params: table?, callback: fun(err?: lsp.ResponseError, result: any))>?

  ---@class test.functional.plugin.lsp.testutil._create_server.Message
  ---@field method vim.lsp.protocol.Method.ClientToServer
  ---@field params table?

  ---@class test.functional.plugin.lsp.testutil._create_server.Server
  ---@field messages test.functional.plugin.lsp.testutil._create_server.Message[]
  ---@field cmd fun(dispatchers: vim.lsp.rpc.Dispatchers, config: vim.lsp.ClientConfig): vim.lsp.rpc.Client

  ---@param opts test.functional.plugin.lsp.testutil._create_server.Opts?
  ---@return test.functional.plugin.lsp.testutil._create_server.Server
  function _G._create_server(opts)
    opts = opts or {}
    local server = {}
    server.messages = {}

    function server.cmd(dispatchers, _config)
      local closing = false
      local handlers = opts.handlers or {}
      local srv = {}

      function srv.request(method, params, callback)
        table.insert(server.messages, {
          method = method,
          params = params,
        })
        local handler = handlers[method]
        if handler then
          handler(method, params, callback)
        elseif method == 'initialize' then
          callback(nil, {
            capabilities = opts.capabilities or {},
          })
        elseif method == 'shutdown' then
          callback(nil, nil)
        end
        local request_id = #server.messages
        return true, request_id
      end

      function srv.notify(method, params)
        table.insert(server.messages, {
          method = method,
          params = params,
        })
        if method == 'exit' then
          dispatchers.on_exit(0, 15)
        end
        return true
      end

      function srv.is_closing()
        return closing
      end

      function srv.terminate()
        closing = true
        dispatchers.on_exit(0, 15)
      end

      return srv
    end

    return server
  end
end

-- Fake LSP server.
M.fake_lsp_code = 'test/functional/fixtures/fake-lsp-server.lua'
M.fake_lsp_logfile = 'Xtest-fake-lsp.log'

local function fake_lsp_server_setup(test_name, timeout_ms, options, settings)
  exec_lua(function(fake_lsp_code, fake_lsp_logfile, timeout)
    options = options or {}
    settings = settings or {}
    _G.lsp = require('vim.lsp')
    _G.TEST_RPC_CLIENT_ID = _G.lsp.start_client {
      cmd_env = {
        NVIM_LOG_FILE = fake_lsp_logfile,
        NVIM_LUA_NOTRACK = '1',
        NVIM_APPNAME = 'nvim_lsp_test',
      },
      cmd = {
        vim.v.progpath,
        '-l',
        fake_lsp_code,
        test_name,
        tostring(timeout),
      },
      handlers = setmetatable({}, {
        __index = function(_t, _method)
          return function(...)
            return vim.rpcrequest(1, 'handler', ...)
          end
        end,
      }),
      workspace_folders = {
        {
          uri = 'file://' .. vim.uv.cwd(),
          name = 'test_folder',
        },
      },
      before_init = function(_params, _config)
        vim.schedule(function()
          vim.rpcrequest(1, 'setup')
        end)
      end,
      on_init = function(client, result)
        _G.TEST_RPC_CLIENT = client
        vim.rpcrequest(1, 'init', result)
      end,
      flags = {
        allow_incremental_sync = options.allow_incremental_sync or false,
        debounce_text_changes = options.debounce_text_changes or 0,
      },
      settings = settings,
      on_exit = function(...)
        vim.rpcnotify(1, 'exit', ...)
      end,
    }
  end, M.fake_lsp_code, M.fake_lsp_logfile, timeout_ms or 1e3)
end

--- @class test.lsp.Config
--- @field test_name string
--- @field timeout_ms? integer
--- @field options? table
--- @field settings? table
---
--- @field on_setup? fun()
--- @field on_init? fun(client: vim.lsp.Client, ...)
--- @field on_handler? fun(...)
--- @field on_exit? fun(code: integer, signal: integer)

--- @param config test.lsp.Config
function M.test_rpc_server(config)
  if config.test_name then
    M.clear_notrace()
    fake_lsp_server_setup(
      config.test_name,
      config.timeout_ms or 1e3,
      config.options,
      config.settings
    )
  end
  local client = setmetatable({}, {
    __index = function(t, name)
      -- Workaround for not being able to yield() inside __index for Lua 5.1 :(
      -- Otherwise I would just return the value here.
      return function(arg1, ...)
        local ismethod = arg1 == t
        return exec_lua(function(...)
          local client = _G.TEST_RPC_CLIENT
          if type(client[name]) == 'function' then
            return client[name](ismethod and client or arg1, ...)
          end
          return client[name]
        end, ...)
      end
    end,
  })
  --- @type integer, integer
  local code, signal
  local busy = 0
  local exited = false
  local function on_request(method, args)
    busy = busy + 1
    if method == 'setup' and config.on_setup then
      config.on_setup()
    end
    if method == 'init' and config.on_init then
      config.on_init(client, 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)
      exited = true
      if busy == 0 then
        stop()
      end
    end
    return NIL
  end
  --  TODO specify timeout?
  --  run(on_request, on_notify, nil, 1000)
  run(on_request, on_notify, nil)
  if config.on_exit then
    config.on_exit(code, signal)
  end
  stop()
  if config.test_name then
    api.nvim_exec_autocmds('VimLeavePre', { modeline = false })
  end
end

return M