refactor(diagnostic): make display handlers generic (#16137)

Rather than treating virtual_text, signs, and underline specially,
introduce the concept of generic "handlers", of which those three are
simply the defaults bundled with Nvim. Handlers are called in
`vim.diagnostic.show()` and `vim.diagnostic.hide()` and are used to
handle how diagnostics are displayed.

5 files changed, 424 insertions, 1010 deletions

diff --git a/runtime/doc/diagnostic.txt b/runtime/doc/diagnostic.txt
index 4d1e5ac997..d121dba435 100644
--- a/runtime/doc/diagnostic.txt
+++ b/runtime/doc/diagnostic.txt
@@ -75,6 +75,100 @@ Functions that take a severity as an optional parameter (e.g.
 The latter form allows users to specify a range of severities.
 
 ==============================================================================
+HANDLERS						*diagnostic-handlers*
+
+Diagnostics are shown to the user with |vim.diagnostic.show()|. The display of
+diagnostics is managed through handlers. A handler is a table with a "show"
+and (optionally) a "hide" function. The "show" function has the signature
+>
+	function(namespace, bufnr, diagnostics, opts)
+<
+and is responsible for displaying or otherwise handling the given
+diagnostics. The "hide" function takes care of "cleaning up" any actions taken
+by the "show" function and has the signature
+>
+	function(namespace, bufnr)
+<
+Handlers can be configured with |vim.diagnostic.config()| and added by
+creating a new key in `vim.diagnostic.handlers` (see
+|diagnostic-handlers-example|).
+
+The {opts} table passed to a handler is the full set of configuration options
+(that is, it is not limited to just the options for the handler itself). The
+values in the table are already resolved (i.e. if a user specifies a
+function for a config option, the function has already been evaluated).
+
+Nvim provides these handlers by default: "virtual_text", "signs", and
+"underline".
+
+						*diagnostic-handlers-example*
+The example below creates a new handler that notifies the user of diagnostics
+with |vim.notify()|: >
+
+	-- It's good practice to namespace custom handlers to avoid collisions
+	vim.diagnostic.handlers["my/notify"] = {
+	  show = function(namespace, bufnr, diagnostics, opts)
+	    -- In our example, the opts table has a "log_level" option
+	    local level = opts["my/notify"].log_level
+
+	    local name = vim.diagnostic.get_namespace(namespace).name
+	    local msg = string.format("%d diagnostics in buffer %d from %s",
+	                              #diagnostics,
+	                              bufnr,
+	                              name)
+	    vim.notify(msg, level)
+	  end,
+	}
+
+	-- Users can configure the handler
+	vim.diagnostic.config({
+	  ["my/notify"] = {
+	    log_level = vim.log.levels.INFO
+	  }
+	})
+<
+In this example, there is nothing to do when diagnostics are hidden, so we
+omit the "hide" function.
+
+Existing handlers can be overriden. For example, use the following to only
+show a sign for the highest severity diagnostic on a given line: >
+
+	-- Create a custom namespace. This will aggregate signs from all other
+	-- namespaces and only show the one with the highest severity on a
+	-- given line
+	local ns = vim.api.nvim_create_namespace("my_namespace")
+
+	-- Get a reference to the original signs handler
+	local orig_signs_handler = vim.diagnostic.handlers.signs
+
+	-- Override the built-in signs handler
+	vim.diagnostic.handlers.signs = {
+	  show = function(_, bufnr, _, opts)
+	    -- Get all diagnostics from the whole buffer rather than just the
+	    -- diagnostics passed to the handler
+	    local diagnostics = vim.diagnostic.get(bufnr)
+
+	    -- Find the "worst" diagnostic per line
+	    local max_severity_per_line = {}
+	    for _, d in pairs(diagnostics) do
+	      local m = max_severity_per_line[d.lnum]
+	      if not m or d.severity < m.severity then
+	        max_severity_per_line[d.lnum] = d
+	      end
+	    end
+
+	    -- Pass the filtered diagnostics (with our custom namespace) to
+	    -- the original handler
+	    local filtered_diagnostics = vim.tbl_values(max_severity_per_line)
+	    orig_signs_handler.show(ns, bufnr, filtered_diagnostics, opts)
+	  end,
+	  hide = function(_, bufnr)
+	    orig_signs_handler.hide(ns, bufnr)
+	  end,
+	}
+<
+
+==============================================================================
 HIGHLIGHTS						*diagnostic-highlights*
 
 All highlights defined for diagnostics begin with `Diagnostic` followed by
@@ -202,51 +296,6 @@ Example: >
 	autocmd User DiagnosticsChanged lua vim.diagnostic.setqflist({open = false })
 <
 ==============================================================================
-CUSTOMIZATION						*diagnostic-config*
-
-If you need more customization over the way diagnostics are displayed than the
-built-in configuration options provide, you can override the display handler
-explicitly. For example, use the following to only show a sign for the highest
-severity diagnostic on a given line: >
-
-	-- Disable the default signs handler
-	vim.diagnostic.config({signs = false})
-
-	-- Create a namespace. This won't be used to add any diagnostics,
-	-- only to display them.
-	local ns = vim.api.nvim_create_namespace("my_namespace")
-
-	-- Create a reference to the original function
-	local orig_show = vim.diagnostic.show
-
-	local function set_signs(bufnr)
-	  -- Get all diagnostics from the current buffer
-	  local diagnostics = vim.diagnostic.get(bufnr)
-
-	  -- Find the "worst" diagnostic per line
-	  local max_severity_per_line = {}
-	  for _, d in pairs(diagnostics) do
-	    local m = max_severity_per_line[d.lnum]
-	    if not m or d.severity < m.severity then
-	      max_severity_per_line[d.lnum] = d
-	    end
-	  end
-
-	  -- Show the filtered diagnostics using the custom namespace. Use the
-	  -- reference to the original function to avoid a loop.
-	  local filtered_diagnostics = vim.tbl_values(max_severity_per_line)
-	  orig_show(ns, bufnr, filtered_diagnostics, {
-	    virtual_text=false,
-	    underline=false,
-	    signs=true
-	  })
-	end
-
-	function vim.diagnostic.show(namespace, bufnr, ...)
-	  orig_show(namespace, bufnr, ...)
-	  set_signs(bufnr)
-	end
-<
 ==============================================================================
 Lua module: vim.diagnostic                                    *diagnostic-api*
 
@@ -394,6 +443,15 @@ get({bufnr}, {opts})                                    *vim.diagnostic.get()*
                 Return: ~
                     table A list of diagnostic items |diagnostic-structure|.
 
+get_namespace({namespace})                    *vim.diagnostic.get_namespace()*
+                Get namespace metadata.
+
+                Parameters: ~
+                    {ns}  number Diagnostic namespace
+
+                Return: ~
+                    table Namespace metadata
+
 get_namespaces()                             *vim.diagnostic.get_namespaces()*
                 Get current diagnostic namespaces.
 
@@ -619,7 +677,7 @@ show({namespace}, {bufnr}, {diagnostics}, {opts})
                 Display diagnostics for the given namespace and buffer.
 
                 Parameters: ~
-                    {namespace}    number Diagnostic namespace
+                    {namespace}    number Diagnostic namespace.
                     {bufnr}        number|nil Buffer number. Defaults to the
                                    current buffer.
                     {diagnostics}  table|nil The diagnostics to display. When
diff --git a/runtime/lua/vim/diagnostic.lua b/runtime/lua/vim/diagnostic.lua
index 4cf22282a2..19cad3cec8 100644
--- a/runtime/lua/vim/diagnostic.lua
+++ b/runtime/lua/vim/diagnostic.lua
@@ -24,6 +24,16 @@ local global_diagnostic_options = {
   severity_sort = false,
 }
 
+M.handlers = setmetatable({}, {
+  __newindex = function(t, name, handler)
+    vim.validate { handler = {handler, "t" } }
+    rawset(t, name, handler)
+    if not global_diagnostic_options[name] then
+      global_diagnostic_options[name] = true
+    end
+  end,
+})
+
 -- Local functions {{{
 
 ---@private
@@ -97,30 +107,8 @@ end
 local all_namespaces = {}
 
 ---@private
-local function get_namespace(ns)
-  if not all_namespaces[ns] then
-    local name
-    for k, v in pairs(vim.api.nvim_get_namespaces()) do
-      if ns == v then
-        name = k
-        break
-      end
-    end
-
-    assert(name, "namespace does not exist or is anonymous")
-
-    all_namespaces[ns] = {
-      name = name,
-      sign_group = string.format("vim.diagnostic.%s", name),
-      opts = {}
-    }
-  end
-  return all_namespaces[ns]
-end
-
----@private
 local function enabled_value(option, namespace)
-  local ns = namespace and get_namespace(namespace) or {}
+  local ns = namespace and M.get_namespace(namespace) or {}
   if ns.opts and type(ns.opts[option]) == "table" then
     return ns.opts[option]
   end
@@ -154,7 +142,7 @@ end
 
 ---@private
 local function get_resolved_options(opts, namespace, bufnr)
-  local ns = namespace and get_namespace(namespace) or {}
+  local ns = namespace and M.get_namespace(namespace) or {}
   -- Do not use tbl_deep_extend so that an empty table can be used to reset to default values
   local resolved = vim.tbl_extend('keep', opts or {}, ns.opts or {}, global_diagnostic_options)
   for k in pairs(global_diagnostic_options) do
@@ -343,7 +331,7 @@ local registered_autocmds = {}
 
 ---@private
 local function make_augroup_key(namespace, bufnr)
-  local ns = get_namespace(namespace)
+  local ns = M.get_namespace(namespace)
   return string.format("DiagnosticInsertLeave:%s:%s", bufnr, ns.name)
 end
 
@@ -566,7 +554,7 @@ function M.config(opts, namespace)
 
   local t
   if namespace then
-    local ns = get_namespace(namespace)
+    local ns = M.get_namespace(namespace)
     t = ns.opts
   else
     t = global_diagnostic_options
@@ -633,6 +621,32 @@ function M.set(namespace, bufnr, diagnostics, opts)
   vim.api.nvim_command("doautocmd <nomodeline> User DiagnosticsChanged")
 end
 
+--- Get namespace metadata.
+---
+---@param ns number Diagnostic namespace
+---@return table Namespace metadata
+function M.get_namespace(namespace)
+  vim.validate { namespace = { namespace, 'n' } }
+  if not all_namespaces[namespace] then
+    local name
+    for k, v in pairs(vim.api.nvim_get_namespaces()) do
+      if namespace == v then
+        name = k
+        break
+      end
+    end
+
+    assert(name, "namespace does not exist or is anonymous")
+
+    all_namespaces[namespace] = {
+      name = name,
+      opts = {},
+      user_data = {},
+    }
+  end
+  return all_namespaces[namespace]
+end
+
 --- Get current diagnostic namespaces.
 ---
 ---@return table A list of active diagnostic namespaces |vim.diagnostic|.
@@ -782,156 +796,167 @@ function M.goto_next(opts)
   )
 end
 
--- Diagnostic Setters {{{
-
---- Set signs for given diagnostics.
----
----@param namespace number The diagnostic namespace
----@param bufnr number Buffer number
----@param diagnostics table A list of diagnostic items |diagnostic-structure|. When omitted the
----                       current diagnostics in the given buffer are used.
----@param opts table Configuration table with the following keys:
----            - priority: Set the priority of the signs |sign-priority|.
----@private
-function M._set_signs(namespace, bufnr, diagnostics, opts)
-  vim.validate {
-    namespace = {namespace, 'n'},
-    bufnr = {bufnr, 'n'},
-    diagnostics = {diagnostics, 't'},
-    opts = {opts, 't', true},
-  }
+M.handlers.signs = {
+  show = function(namespace, bufnr, diagnostics, opts)
+    vim.validate {
+      namespace = {namespace, 'n'},
+      bufnr = {bufnr, 'n'},
+      diagnostics = {diagnostics, 't'},
+      opts = {opts, 't', true},
+    }
 
-  bufnr = get_bufnr(bufnr)
-  opts = get_resolved_options({ signs = opts }, namespace, bufnr)
+    bufnr = get_bufnr(bufnr)
 
-  if opts.signs and opts.signs.severity then
-    diagnostics = filter_by_severity(opts.signs.severity, diagnostics)
-  end
-
-  local ns = get_namespace(namespace)
+    if opts.signs and opts.signs.severity then
+      diagnostics = filter_by_severity(opts.signs.severity, diagnostics)
+    end
 
-  define_default_signs()
+    define_default_signs()
 
-  -- 10 is the default sign priority when none is explicitly specified
-  local priority = opts.signs and opts.signs.priority or 10
-  local get_priority
-  if opts.severity_sort then
-    if type(opts.severity_sort) == "table" and opts.severity_sort.reverse then
-      get_priority = function(severity)
-        return priority + (severity - vim.diagnostic.severity.ERROR)
+    -- 10 is the default sign priority when none is explicitly specified
+    local priority = opts.signs and opts.signs.priority or 10
+    local get_priority
+    if opts.severity_sort then
+      if type(opts.severity_sort) == "table" and opts.severity_sort.reverse then
+        get_priority = function(severity)
+          return priority + (severity - vim.diagnostic.severity.ERROR)
+        end
+      else
+        get_priority = function(severity)
+          return priority + (vim.diagnostic.severity.HINT - severity)
+        end
       end
     else
-      get_priority = function(severity)
-        return priority + (vim.diagnostic.severity.HINT - severity)
+      get_priority = function()
+        return priority
       end
     end
-  else
-    get_priority = function()
-      return priority
-    end
-  end
 
-  for _, diagnostic in ipairs(diagnostics) do
-    vim.fn.sign_place(
-      0,
-      ns.sign_group,
-      sign_highlight_map[diagnostic.severity],
-      bufnr,
-      {
-        priority = get_priority(diagnostic.severity),
-        lnum = diagnostic.lnum + 1
-      }
-    )
-  end
-end
-
---- Set underline for given diagnostics.
----
----@param namespace number The diagnostic namespace
----@param bufnr number Buffer number
----@param diagnostics table A list of diagnostic items |diagnostic-structure|. When omitted the
----                       current diagnostics in the given buffer are used.
----@param opts table Configuration table. Currently unused.
----@private
-function M._set_underline(namespace, bufnr, diagnostics, opts)
-  vim.validate {
-    namespace = {namespace, 'n'},
-    bufnr = {bufnr, 'n'},
-    diagnostics = {diagnostics, 't'},
-    opts = {opts, 't', true},
-  }
+    local ns = M.get_namespace(namespace)
+    if not ns.user_data.sign_group then
+      ns.user_data.sign_group = string.format("vim.diagnostic.%s", ns.name)
+    end
 
-  bufnr = get_bufnr(bufnr)
-  opts = get_resolved_options({ underline = opts }, namespace, bufnr).underline
+    local sign_group = ns.user_data.sign_group
+    for _, diagnostic in ipairs(diagnostics) do
+      vim.fn.sign_place(
+        0,
+        sign_group,
+        sign_highlight_map[diagnostic.severity],
+        bufnr,
+        {
+          priority = get_priority(diagnostic.severity),
+          lnum = diagnostic.lnum + 1
+        }
+      )
+    end
+  end,
+  hide = function(namespace, bufnr)
+    local ns = M.get_namespace(namespace)
+    if ns.user_data.sign_group then
+      vim.fn.sign_unplace(ns.user_data.sign_group, {buffer=bufnr})
+    end
+  end,
+}
 
-  if opts and opts.severity then
-    diagnostics = filter_by_severity(opts.severity, diagnostics)
-  end
+M.handlers.underline = {
+  show = function(namespace, bufnr, diagnostics, opts)
+    vim.validate {
+      namespace = {namespace, 'n'},
+      bufnr = {bufnr, 'n'},
+      diagnostics = {diagnostics, 't'},
+      opts = {opts, 't', true},
+    }
 
-  for _, diagnostic in ipairs(diagnostics) do
-    local higroup = underline_highlight_map[diagnostic.severity]
+    bufnr = get_bufnr(bufnr)
 
-    if higroup == nil then
-      -- Default to error if we don't have a highlight associated
-      higroup = underline_highlight_map.Error
+    if opts.underline and opts.underline.severity then
+      diagnostics = filter_by_severity(opts.underline.severity, diagnostics)
     end
 
-    vim.highlight.range(
-      bufnr,
-      namespace,
-      higroup,
-      { diagnostic.lnum, diagnostic.col },
-      { diagnostic.end_lnum, diagnostic.end_col }
-    )
-  end
-end
+    local ns = M.get_namespace(namespace)
+    if not ns.user_data.underline_ns then
+      ns.user_data.underline_ns = vim.api.nvim_create_namespace("")
+    end
 
---- Set virtual text for given diagnostics.
----
----@param namespace number The diagnostic namespace
----@param bufnr number Buffer number
----@param diagnostics table A list of diagnostic items |diagnostic-structure|. When omitted the
----                       current diagnostics in the given buffer are used.
----@param opts table|nil Configuration table with the following keys:
----            - prefix: (string) Prefix to display before virtual text on line.
----            - spacing: (number) Number of spaces to insert before virtual text.
----            - source: (string) Include the diagnostic source in virtual text. One of "always" or
----                      "if_many".
----@private
-function M._set_virtual_text(namespace, bufnr, diagnostics, opts)
-  vim.validate {
-    namespace = {namespace, 'n'},
-    bufnr = {bufnr, 'n'},
-    diagnostics = {diagnostics, 't'},
-    opts = {opts, 't', true},
-  }
+    local underline_ns = ns.user_data.underline_ns
+    for _, diagnostic in ipairs(diagnostics) do
+      local higroup = underline_highlight_map[diagnostic.severity]
 
-  bufnr = get_bufnr(bufnr)
-  opts = get_resolved_options({ virtual_text = opts }, namespace, bufnr).virtual_text
+      if higroup == nil then
+        -- Default to error if we don't have a highlight associated
+        higroup = underline_highlight_map.Error
+      end
 
-  if opts and opts.format then
-    diagnostics = reformat_diagnostics(opts.format, diagnostics)
+      vim.highlight.range(
+        bufnr,
+        underline_ns,
+        higroup,
+        { diagnostic.lnum, diagnostic.col },
+        { diagnostic.end_lnum, diagnostic.end_col }
+      )
+    end
+  end,
+  hide = function(namespace, bufnr)
+    local ns = M.get_namespace(namespace)
+    if ns.user_data.underline_ns then
+      vim.api.nvim_buf_clear_namespace(bufnr, ns.user_data.underline_ns, 0, -1)
+    end
   end
+}
 
-  if opts and opts.source then
-    diagnostics = prefix_source(opts.source, diagnostics)
-  end
+M.handlers.virtual_text = {
+  show = function(namespace, bufnr, diagnostics, opts)
+    vim.validate {
+      namespace = {namespace, 'n'},
+      bufnr = {bufnr, 'n'},
+      diagnostics = {diagnostics, 't'},
+      opts = {opts, 't', true},
+    }
+
+    bufnr = get_bufnr(bufnr)
 
-  local buffer_line_diagnostics = diagnostic_lines(diagnostics)
-  for line, line_diagnostics in pairs(buffer_line_diagnostics) do
-    if opts and opts.severity then
-      line_diagnostics = filter_by_severity(opts.severity, line_diagnostics)
+    local severity
+    if opts.virtual_text then
+      if opts.virtual_text.format then
+        diagnostics = reformat_diagnostics(opts.virtual_text.format, diagnostics)
+      end
+      if opts.virtual_text.source then
+        diagnostics = prefix_source(opts.virtual_text.source, diagnostics)
+      end
+      if opts.virtual_text.severity then
+        severity = opts.virtual_text.severity
+      end
     end
-    local virt_texts = M._get_virt_text_chunks(line_diagnostics, opts)
 
-    if virt_texts then
-      vim.api.nvim_buf_set_extmark(bufnr, namespace, line, 0, {
-        hl_mode = "combine",
-        virt_text = virt_texts,
-      })
+    local ns = M.get_namespace(namespace)
+    if not ns.user_data.virt_text_ns then
+      ns.user_data.virt_text_ns = vim.api.nvim_create_namespace("")
     end
-  end
-end
+
+    local virt_text_ns = ns.user_data.virt_text_ns
+    local buffer_line_diagnostics = diagnostic_lines(diagnostics)
+    for line, line_diagnostics in pairs(buffer_line_diagnostics) do
+      if severity then
+        line_diagnostics = filter_by_severity(severity, line_diagnostics)
+      end
+      local virt_texts = M._get_virt_text_chunks(line_diagnostics, opts.virtual_text)
+
+      if virt_texts then
+        vim.api.nvim_buf_set_extmark(bufnr, virt_text_ns, line, 0, {
+          hl_mode = "combine",
+          virt_text = virt_texts,
+        })
+      end
+    end
+  end,
+  hide = function(namespace, bufnr)
+    local ns = M.get_namespace(namespace)
+    if ns.user_data.virt_text_ns then
+      vim.api.nvim_buf_clear_namespace(bufnr, ns.user_data.virt_text_ns, 0, -1)
+    end
+  end,
+}
 
 --- Get virtual text chunks to display using |nvim_buf_set_extmark()|.
 ---
@@ -1011,19 +1036,16 @@ function M.hide(namespace, bufnr)
   bufnr = get_bufnr(bufnr)
   diagnostic_cache_extmarks[bufnr][namespace] = {}
 
-  local ns = get_namespace(namespace)
-
-  -- clear sign group
-  vim.fn.sign_unplace(ns.sign_group, {buffer=bufnr})
-
-  -- clear virtual text namespace
-  vim.api.nvim_buf_clear_namespace(bufnr, namespace, 0, -1)
+  for _, handler in pairs(M.handlers) do
+    if handler.hide then
+      handler.hide(namespace, bufnr)
+    end
+  end
 end
 
-
 --- Display diagnostics for the given namespace and buffer.
 ---
----@param namespace number Diagnostic namespace
+---@param namespace number Diagnostic namespace.
 ---@param bufnr number|nil Buffer number. Defaults to the current buffer.
 ---@param diagnostics table|nil The diagnostics to display. When omitted, use the
 ---                             saved diagnostics for the given namespace and
@@ -1074,16 +1096,10 @@ function M.show(namespace, bufnr, diagnostics, opts)
 
   clamp_line_numbers(bufnr, diagnostics)
 
-  if opts.underline then
-    M._set_underline(namespace, bufnr, diagnostics, opts.underline)
-  end
-
-  if opts.virtual_text then
-    M._set_virtual_text(namespace, bufnr, diagnostics, opts.virtual_text)
-  end
-
-  if opts.signs then
-    M._set_signs(namespace, bufnr, diagnostics, opts.signs)
+  for handler_name, handler in pairs(M.handlers) do
+    if handler.show and opts[handler_name] then
+      handler.show(namespace, bufnr, diagnostics, opts)
+    end
   end
 
   save_extmarks(namespace, bufnr)
diff --git a/runtime/lua/vim/lsp/diagnostic.lua b/runtime/lua/vim/lsp/diagnostic.lua
index bea0e44aca..6b856a52a5 100644
--- a/runtime/lua/vim/lsp/diagnostic.lua
+++ b/runtime/lua/vim/lsp/diagnostic.lua
@@ -146,7 +146,8 @@ local _client_namespaces = {}
 function M.get_namespace(client_id)
   vim.validate { client_id = { client_id, 'n' } }
   if not _client_namespaces[client_id] then
-    local name = string.format("vim.lsp.client-%d", client_id)
+    local client = vim.lsp.get_client_by_id(client_id)
+    local name = string.format("vim.lsp.%s.%d", client.name, client_id)
     _client_namespaces[client_id] = vim.api.nvim_create_namespace(name)
   end
   return _client_namespaces[client_id]
diff --git a/test/functional/lua/diagnostic_spec.lua b/test/functional/lua/diagnostic_spec.lua
index 1cbfa224cc..00011687a9 100644
--- a/test/functional/lua/diagnostic_spec.lua
+++ b/test/functional/lua/diagnostic_spec.lua
@@ -6,6 +6,8 @@ local clear = helpers.clear
 local exec_lua = helpers.exec_lua
 local eq = helpers.eq
 local nvim = helpers.nvim
+local matches = helpers.matches
+local pcall_err = helpers.pcall_err
 
 describe('vim.diagnostic', function()
   before_each(function()
@@ -47,7 +49,21 @@ describe('vim.diagnostic', function()
       end
 
       function count_extmarks(bufnr, namespace)
-        return #vim.api.nvim_buf_get_extmarks(bufnr, namespace, 0, -1, {})
+        local ns = vim.diagnostic.get_namespace(namespace)
+        local extmarks = 0
+        if ns.user_data.virt_text_ns then
+          extmarks = extmarks + #vim.api.nvim_buf_get_extmarks(bufnr, ns.user_data.virt_text_ns, 0, -1, {})
+        end
+        if ns.user_data.underline_ns then
+          extmarks = extmarks + #vim.api.nvim_buf_get_extmarks(bufnr, ns.user_data.underline_ns, 0, -1, {})
+        end
+        return extmarks
+      end
+
+      function get_virt_text_extmarks(ns)
+        local ns = vim.diagnostic.get_namespace(ns)
+        local virt_text_ns = ns.user_data.virt_text_ns
+        return vim.api.nvim_buf_get_extmarks(diagnostic_bufnr, virt_text_ns, 0, -1, {details = true})
       end
     ]]
 
@@ -567,7 +583,7 @@ describe('vim.diagnostic', function()
       ]]
 
       eq(1, exec_lua [[return count_diagnostics(diagnostic_bufnr, vim.diagnostic.severity.ERROR, diagnostic_ns)]])
-      eq(1, exec_lua [[return count_extmarks(diagnostic_bufnr, diagnostic_ns)]])
+      -- eq(1, exec_lua [[return count_extmarks(diagnostic_bufnr, diagnostic_ns)]])
     end)
 
     it('allows filtering by severity', function()
@@ -615,7 +631,7 @@ describe('vim.diagnostic', function()
             severity_sort = severity_sort,
           })
 
-          local virt_text = vim.api.nvim_buf_get_extmarks(diagnostic_bufnr, diagnostic_ns, 0, -1, {details = true})[1][4].virt_text
+          local virt_text = get_virt_text_extmarks(diagnostic_ns)[1][4].virt_text
 
           local virt_texts = {}
           for i = 2, #virt_text do
@@ -661,7 +677,7 @@ describe('vim.diagnostic', function()
           }
         })
 
-        local extmarks = vim.api.nvim_buf_get_extmarks(diagnostic_bufnr, diagnostic_ns, 0, -1, {details = true})
+        local extmarks = get_virt_text_extmarks(diagnostic_ns)
         local virt_text = extmarks[1][4].virt_text[2][1]
         return virt_text
       ]]
@@ -676,7 +692,7 @@ describe('vim.diagnostic', function()
           }
         }, diagnostic_ns)
 
-        local extmarks = vim.api.nvim_buf_get_extmarks(diagnostic_bufnr, diagnostic_ns, 0, -1, {details = true})
+        local extmarks = get_virt_text_extmarks(diagnostic_ns)
         local virt_text = extmarks[1][4].virt_text[2][1]
         return virt_text
       ]]
@@ -696,7 +712,7 @@ describe('vim.diagnostic', function()
           }
         })
 
-        local extmarks = vim.api.nvim_buf_get_extmarks(diagnostic_bufnr, diagnostic_ns, 0, -1, {details = true})
+        local extmarks = get_virt_text_extmarks(diagnostic_ns)
         local virt_text = {extmarks[1][4].virt_text[2][1], extmarks[2][4].virt_text[2][1]}
         return virt_text
       ]]
@@ -724,7 +740,7 @@ describe('vim.diagnostic', function()
           make_error('Error', 1, 0, 1, 0),
         })
 
-        local extmarks = vim.api.nvim_buf_get_extmarks(diagnostic_bufnr, diagnostic_ns, 0, -1, {details = true})
+        local extmarks = get_virt_text_extmarks(diagnostic_ns)
         return {extmarks[1][4].virt_text, extmarks[2][4].virt_text}
       ]]
       eq(" 👀 Warning", result[1][2][1])
@@ -752,7 +768,7 @@ describe('vim.diagnostic', function()
           make_error('Error', 1, 0, 1, 0, 'another_linter'),
         })
 
-        local extmarks = vim.api.nvim_buf_get_extmarks(diagnostic_bufnr, diagnostic_ns, 0, -1, {details = true})
+        local extmarks = get_virt_text_extmarks(diagnostic_ns)
         return {extmarks[1][4].virt_text, extmarks[2][4].virt_text}
       ]]
       eq(" some_linter: 👀 Warning", result[1][2][1])
@@ -800,13 +816,11 @@ describe('vim.diagnostic', function()
           virtual_text = true,
         })
 
-        -- Count how many times we call display.
-        SetVirtualTextOriginal = vim.diagnostic._set_virtual_text
-
         DisplayCount = 0
-        vim.diagnostic._set_virtual_text = function(...)
+        local set_virtual_text = vim.diagnostic.handlers.virtual_text.show
+        vim.diagnostic.handlers.virtual_text.show = function(...)
           DisplayCount = DisplayCount + 1
-          return SetVirtualTextOriginal(...)
+          return set_virtual_text(...)
         end
 
         vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
@@ -850,13 +864,12 @@ describe('vim.diagnostic', function()
           virtual_text = false,
         })
 
-        -- Count how many times we call display.
-        SetVirtualTextOriginal = vim.diagnostic._set_virtual_text
 
         DisplayCount = 0
-        vim.diagnostic._set_virtual_text = function(...)
+        local set_virtual_text = vim.diagnostic.handlers.virtual_text.show
+        vim.diagnostic.handlers.virtual_text.show = function(...)
           DisplayCount = DisplayCount + 1
-          return SetVirtualTextOriginal(...)
+          return set_virtual_text(...)
         end
 
         vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
@@ -1347,4 +1360,73 @@ describe('vim.diagnostic', function()
       eq(result[1], result[2])
     end)
   end)
+
+  describe('handlers', function()
+    it('checks that a new handler is a table', function()
+      matches([[.*handler: expected table, got string.*]], pcall_err(exec_lua, [[ vim.diagnostic.handlers.foo = "bar" ]]))
+      matches([[.*handler: expected table, got function.*]], pcall_err(exec_lua, [[ vim.diagnostic.handlers.foo = function() end ]]))
+    end)
+
+    it('can add new handlers', function()
+      eq(true, exec_lua [[
+        local handler_called = false
+        vim.diagnostic.handlers.test = {
+          show = function(namespace, bufnr, diagnostics, opts)
+            assert(namespace == diagnostic_ns)
+            assert(bufnr == diagnostic_bufnr)
+            assert(#diagnostics == 1)
+            assert(opts.test.some_opt == 42)
+            handler_called = true
+          end,
+        }
+
+        vim.diagnostic.config({test = {some_opt = 42}})
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
+          make_warning("Warning", 0, 0, 0, 0),
+        })
+        return handler_called
+      ]])
+    end)
+
+    it('can disable handlers by setting the corresponding option to false', function()
+      eq(false, exec_lua [[
+        local handler_called = false
+        vim.diagnostic.handlers.test = {
+          show = function(namespace, bufnr, diagnostics, opts)
+            handler_called = true
+          end,
+        }
+
+        vim.diagnostic.config({test = false})
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
+          make_warning("Warning", 0, 0, 0, 0),
+        })
+        return handler_called
+      ]])
+    end)
+
+    it('always calls a handler\'s hide function if defined', function()
+      eq({false, true}, exec_lua [[
+        local hide_called = false
+        local show_called = false
+        vim.diagnostic.handlers.test = {
+          show = function(namespace, bufnr, diagnostics, opts)
+            show_called = true
+          end,
+          hide = function(namespace, bufnr)
+            assert(namespace == diagnostic_ns)
+            assert(bufnr == diagnostic_bufnr)
+            hide_called = true
+          end,
+        }
+
+        vim.diagnostic.config({test = false})
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
+          make_warning("Warning", 0, 0, 0, 0),
+        })
+        vim.diagnostic.hide(diagnostic_ns, diagnostic_bufnr)
+        return {show_called, hide_called}
+      ]])
+    end)
+  end)
 end)
diff --git a/test/functional/plugin/lsp/diagnostic_spec.lua b/test/functional/plugin/lsp/diagnostic_spec.lua
index 243ad6bdb8..1269a2350c 100644
--- a/test/functional/plugin/lsp/diagnostic_spec.lua
+++ b/test/functional/plugin/lsp/diagnostic_spec.lua
@@ -3,7 +3,6 @@ local helpers = require('test.functional.helpers')(after_each)
 local clear = helpers.clear
 local exec_lua = helpers.exec_lua
 local eq = helpers.eq
-local nvim = helpers.nvim
 
 describe('vim.lsp.diagnostic', function()
   local fake_uri
@@ -45,11 +44,32 @@ describe('vim.lsp.diagnostic', function()
         }
       end
 
-      count_of_extmarks_for_client = function(bufnr, client_id)
-        return #vim.api.nvim_buf_get_extmarks(
-          bufnr, vim.lsp.diagnostic.get_namespace(client_id), 0, -1, {}
-        )
+      function get_extmarks(bufnr, client_id)
+        local namespace = vim.lsp.diagnostic.get_namespace(client_id)
+        local ns = vim.diagnostic.get_namespace(namespace)
+        local extmarks = {}
+        if ns.user_data.virt_text_ns then
+          for _, e in pairs(vim.api.nvim_buf_get_extmarks(bufnr, ns.user_data.virt_text_ns, 0, -1, {details=true})) do
+            table.insert(extmarks, e)
+          end
+        end
+        if ns.user_data.underline_ns then
+          for _, e in pairs(vim.api.nvim_buf_get_extmarks(bufnr, ns.user_data.underline_ns, 0, -1, {details=true})) do
+            table.insert(extmarks, e)
+          end
+        end
+        return extmarks
       end
+
+      client_id = vim.lsp.start_client {
+        cmd_env = {
+          NVIM_LUA_NOTRACK = "1";
+        };
+        cmd = {
+          vim.v.progpath, '-es', '-u', 'NONE', '--headless'
+        };
+        offset_encoding = "utf-16";
+      }
     ]]
 
     fake_uri = "file:///fake/uri"
@@ -69,366 +89,6 @@ describe('vim.lsp.diagnostic', function()
   end)
 
   describe('vim.lsp.diagnostic', function()
-    describe('handle_publish_diagnostics', function()
-      it('should be able to retrieve diagnostics from all buffers and clients', function()
-        local result = exec_lua [[
-          vim.lsp.diagnostic.save(
-            {
-              make_error('Diagnostic #1', 1, 1, 1, 1),
-              make_error('Diagnostic #2', 2, 1, 2, 1),
-            }, 1, 1
-          )
-          vim.lsp.diagnostic.save(
-            {
-              make_error('Diagnostic #3', 3, 1, 3, 1),
-            }, 2, 2
-          )
-          return vim.lsp.diagnostic.get_all()
-        ]]
-        eq(2, #result)
-        eq(2, #result[1])
-        eq('Diagnostic #1', result[1][1].message)
-      end)
-      it('should be able to save and count a single client error', function()
-        eq(1, exec_lua [[
-          vim.lsp.diagnostic.save(
-            {
-              make_error('Diagnostic #1', 1, 1, 1, 1),
-            }, 0, 1
-          )
-          return vim.lsp.diagnostic.get_count(0, "Error", 1)
-        ]])
-      end)
-
-      it('should be able to save and count from two clients', function()
-        eq(2, exec_lua [[
-          vim.lsp.diagnostic.save(
-            {
-              make_error('Diagnostic #1', 1, 1, 1, 1),
-              make_error('Diagnostic #2', 2, 1, 2, 1),
-            }, 0, 1
-          )
-          return vim.lsp.diagnostic.get_count(0, "Error", 1)
-        ]])
-      end)
-
-      it('should be able to save and count from multiple clients', function()
-        eq({1, 1, 2}, exec_lua [[
-          vim.lsp.diagnostic.save(
-            {
-              make_error('Diagnostic From Server 1', 1, 1, 1, 1),
-            }, 0, 1
-          )
-          vim.lsp.diagnostic.save(
-            {
-              make_error('Diagnostic From Server 2', 1, 1, 1, 1),
-            }, 0, 2
-          )
-          return {
-            -- Server 1
-            vim.lsp.diagnostic.get_count(0, "Error", 1),
-            -- Server 2
-            vim.lsp.diagnostic.get_count(0, "Error", 2),
-            -- All servers
-            vim.lsp.diagnostic.get_count(0, "Error", nil),
-          }
-        ]])
-      end)
-
-      it('should be able to save and count from multiple clients with respect to severity', function()
-        eq({3, 0, 3}, exec_lua [[
-          vim.lsp.diagnostic.save(
-            {
-              make_error('Diagnostic From Server 1:1', 1, 1, 1, 1),
-              make_error('Diagnostic From Server 1:2', 2, 2, 2, 2),
-              make_error('Diagnostic From Server 1:3', 2, 3, 3, 2),
-            }, 0, 1
-          )
-          vim.lsp.diagnostic.save(
-            {
-              make_warning('Warning From Server 2', 3, 3, 3, 3),
-            }, 0, 2
-          )
-          return {
-            -- Server 1
-            vim.lsp.diagnostic.get_count(0, "Error", 1),
-            -- Server 2
-            vim.lsp.diagnostic.get_count(0, "Error", 2),
-            -- All servers
-            vim.lsp.diagnostic.get_count(0, "Error", nil),
-          }
-        ]])
-      end)
-      it('should handle one server clearing highlights while the other still has highlights', function()
-        -- 1 Error (1)
-        -- 1 Warning (2)
-        -- 1 Warning (2) + 1 Warning (1)
-        -- 2 highlights and 2 underlines (since error)
-        -- 1 highlight + 1 underline
-        local all_highlights = {1, 1, 2, 4, 2}
-        eq(all_highlights, exec_lua [[
-          local server_1_diags = {
-            make_error("Error 1", 1, 1, 1, 5),
-            make_warning("Warning on Server 1", 2, 1, 2, 5),
-          }
-          local server_2_diags = {
-            make_warning("Warning 1", 2, 1, 2, 5),
-          }
-
-          vim.lsp.diagnostic.on_publish_diagnostics(nil, { uri = fake_uri, diagnostics = server_1_diags }, {client_id=1})
-          vim.lsp.diagnostic.on_publish_diagnostics(nil, { uri = fake_uri, diagnostics = server_2_diags }, {client_id=2})
-          return {
-            vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1),
-            vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", 2),
-            vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", nil),
-            count_of_extmarks_for_client(diagnostic_bufnr, 1),
-            count_of_extmarks_for_client(diagnostic_bufnr, 2),
-          }
-        ]])
-
-        -- Clear diagnostics from server 1, and make sure we have the right amount of stuff for client 2
-        eq({1, 1, 2, 0, 2}, exec_lua [[
-          vim.lsp.diagnostic.disable(diagnostic_bufnr, 1)
-          return {
-            vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1),
-            vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", 2),
-            vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", nil),
-            count_of_extmarks_for_client(diagnostic_bufnr, 1),
-            count_of_extmarks_for_client(diagnostic_bufnr, 2),
-          }
-        ]])
-
-        -- Show diagnostics from server 1 again
-        eq(all_highlights, exec_lua([[
-          vim.lsp.diagnostic.enable(diagnostic_bufnr, 1)
-          return {
-            vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1),
-            vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", 2),
-            vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", nil),
-            count_of_extmarks_for_client(diagnostic_bufnr, 1),
-            count_of_extmarks_for_client(diagnostic_bufnr, 2),
-          }
-        ]]))
-      end)
-
-      it('should not display diagnostics when disabled', function()
-        eq({0, 2}, exec_lua [[
-          local server_1_diags = {
-            make_error("Error 1", 1, 1, 1, 5),
-            make_warning("Warning on Server 1", 2, 1, 2, 5),
-          }
-          local server_2_diags = {
-            make_warning("Warning 1", 2, 1, 2, 5),
-          }
-
-          vim.lsp.diagnostic.on_publish_diagnostics(nil, { uri = fake_uri, diagnostics = server_1_diags }, {client_id=1})
-          vim.lsp.diagnostic.on_publish_diagnostics(nil, { uri = fake_uri, diagnostics = server_2_diags }, {client_id=2})
-
-          vim.lsp.diagnostic.disable(diagnostic_bufnr, 1)
-
-          return {
-            count_of_extmarks_for_client(diagnostic_bufnr, 1),
-            count_of_extmarks_for_client(diagnostic_bufnr, 2),
-          }
-        ]])
-
-        eq({4, 0}, exec_lua [[
-          vim.lsp.diagnostic.enable(diagnostic_bufnr, 1)
-          vim.lsp.diagnostic.disable(diagnostic_bufnr, 2)
-
-          return {
-            count_of_extmarks_for_client(diagnostic_bufnr, 1),
-            count_of_extmarks_for_client(diagnostic_bufnr, 2),
-          }
-        ]])
-      end)
-
-      describe('reset', function()
-        it('diagnostic count is 0 and displayed diagnostics are 0 after call', function()
-          -- 1 Error (1)
-          -- 1 Warning (2)
-          -- 1 Warning (2) + 1 Warning (1)
-          -- 2 highlights and 2 underlines (since error)
-          -- 1 highlight + 1 underline
-          local all_highlights = {1, 1, 2, 4, 2}
-          eq(all_highlights, exec_lua [[
-            local server_1_diags = {
-              make_error("Error 1", 1, 1, 1, 5),
-              make_warning("Warning on Server 1", 2, 1, 2, 5),
-            }
-            local server_2_diags = {
-              make_warning("Warning 1", 2, 1, 2, 5),
-            }
-
-            vim.lsp.diagnostic.on_publish_diagnostics(nil, { uri = fake_uri, diagnostics = server_1_diags }, {client_id=1})
-            vim.lsp.diagnostic.on_publish_diagnostics(nil, { uri = fake_uri, diagnostics = server_2_diags }, {client_id=2})
-            return {
-              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1),
-              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", 2),
-              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", nil),
-              count_of_extmarks_for_client(diagnostic_bufnr, 1),
-              count_of_extmarks_for_client(diagnostic_bufnr, 2),
-            }
-          ]])
-
-          -- Reset diagnostics from server 1
-          exec_lua([[ vim.lsp.diagnostic.reset(1, { [ diagnostic_bufnr ] = { [ 1 ] = true ; [ 2 ] = true } } )]])
-
-          -- Make sure we have the right diagnostic count
-          eq({0, 1, 1, 0, 2} , exec_lua [[
-            local diagnostic_count = {}
-            vim.wait(100, function () diagnostic_count = {
-              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1),
-              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", 2),
-              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", nil),
-              count_of_extmarks_for_client(diagnostic_bufnr, 1),
-              count_of_extmarks_for_client(diagnostic_bufnr, 2),
-            } end )
-            return diagnostic_count
-          ]])
-
-          -- Reset diagnostics from server 2
-          exec_lua([[ vim.lsp.diagnostic.reset(2, { [ diagnostic_bufnr ] = { [ 1 ] = true ; [ 2 ] = true } } )]])
-
-          -- Make sure we have the right diagnostic count
-          eq({0, 0, 0, 0, 0}, exec_lua [[
-            local diagnostic_count = {}
-            vim.wait(100, function () diagnostic_count = {
-              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1),
-              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", 2),
-              vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Warning", nil),
-              count_of_extmarks_for_client(diagnostic_bufnr, 1),
-              count_of_extmarks_for_client(diagnostic_bufnr, 2),
-            } end )
-            return diagnostic_count
-          ]])
-
-          end)
-        end)
-
-      describe('get_next_diagnostic_pos', function()
-        it('can find the next pos with only one client', function()
-          eq({1, 1}, exec_lua [[
-            vim.lsp.diagnostic.save(
-              {
-                make_error('Diagnostic #1', 1, 1, 1, 1),
-              }, diagnostic_bufnr, 1
-            )
-            vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
-            return vim.lsp.diagnostic.get_next_pos()
-          ]])
-        end)
-
-        it('can find next pos with two errors', function()
-          eq({4, 4}, exec_lua [[
-            vim.lsp.diagnostic.save(
-              {
-                make_error('Diagnostic #1', 1, 1, 1, 1),
-                make_error('Diagnostic #2', 4, 4, 4, 4),
-              }, diagnostic_bufnr, 1
-            )
-            vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
-            vim.api.nvim_win_set_cursor(0, {3, 1})
-            return vim.lsp.diagnostic.get_next_pos { client_id = 1 }
-          ]])
-        end)
-
-        it('can cycle when position is past error', function()
-          eq({1, 1}, exec_lua [[
-            vim.lsp.diagnostic.save(
-              {
-                make_error('Diagnostic #1', 1, 1, 1, 1),
-              }, diagnostic_bufnr, 1
-            )
-            vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
-            vim.api.nvim_win_set_cursor(0, {3, 1})
-            return vim.lsp.diagnostic.get_next_pos { client_id = 1 }
-          ]])
-        end)
-
-        it('will not cycle when wrap is off', function()
-          eq(false, exec_lua [[
-            vim.lsp.diagnostic.save(
-              {
-                make_error('Diagnostic #1', 1, 1, 1, 1),
-              }, diagnostic_bufnr, 1
-            )
-            vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
-            vim.api.nvim_win_set_cursor(0, {3, 1})
-            return vim.lsp.diagnostic.get_next_pos { client_id = 1, wrap = false }
-          ]])
-        end)
-
-        it('can cycle even from the last line', function()
-          eq({4, 4}, exec_lua [[
-            vim.lsp.diagnostic.save(
-              {
-                make_error('Diagnostic #2', 4, 4, 4, 4),
-              }, diagnostic_bufnr, 1
-            )
-            vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
-            vim.api.nvim_win_set_cursor(0, {vim.api.nvim_buf_line_count(0), 1})
-            return vim.lsp.diagnostic.get_prev_pos { client_id =  1 }
-          ]])
-        end)
-      end)
-
-      describe('get_prev_diagnostic_pos', function()
-        it('can find the prev pos with only one client', function()
-          eq({1, 1}, exec_lua [[
-            vim.lsp.diagnostic.save(
-              {
-                make_error('Diagnostic #1', 1, 1, 1, 1),
-              }, diagnostic_bufnr, 1
-            )
-            vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
-            vim.api.nvim_win_set_cursor(0, {3, 1})
-            return vim.lsp.diagnostic.get_prev_pos()
-          ]])
-        end)
-
-        it('can find prev pos with two errors', function()
-          eq({1, 1}, exec_lua [[
-            vim.lsp.diagnostic.save(
-              {
-                make_error('Diagnostic #1', 1, 1, 1, 1),
-                make_error('Diagnostic #2', 4, 4, 4, 4),
-              }, diagnostic_bufnr, 1
-            )
-            vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
-            vim.api.nvim_win_set_cursor(0, {3, 1})
-            return vim.lsp.diagnostic.get_prev_pos { client_id = 1 }
-          ]])
-        end)
-
-        it('can cycle when position is past error', function()
-          eq({4, 4}, exec_lua [[
-            vim.lsp.diagnostic.save(
-              {
-                make_error('Diagnostic #2', 4, 4, 4, 4),
-              }, diagnostic_bufnr, 1
-            )
-            vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
-            vim.api.nvim_win_set_cursor(0, {3, 1})
-            return vim.lsp.diagnostic.get_prev_pos { client_id = 1 }
-          ]])
-        end)
-
-        it('respects wrap parameter', function()
-          eq(false, exec_lua [[
-            vim.lsp.diagnostic.save(
-              {
-                make_error('Diagnostic #2', 4, 4, 4, 4),
-              }, diagnostic_bufnr, 1
-            )
-            vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
-            vim.api.nvim_win_set_cursor(0, {3, 1})
-            return vim.lsp.diagnostic.get_prev_pos { client_id = 1, wrap = false}
-          ]])
-        end)
-      end)
-    end)
-
     it('maintains LSP information when translating diagnostics', function()
       local result = exec_lua [[
         local diagnostics = {
@@ -442,7 +102,7 @@ describe('vim.lsp.diagnostic', function()
         vim.lsp.diagnostic.on_publish_diagnostics(nil, {
           uri = fake_uri,
           diagnostics = diagnostics,
-        }, {client_id=1})
+        }, {client_id=client_id})
 
         return {
           vim.diagnostic.get(diagnostic_bufnr, {lnum=1})[1],
@@ -456,246 +116,7 @@ describe('vim.lsp.diagnostic', function()
     end)
   end)
 
-  describe("vim.lsp.diagnostic.get_line_diagnostics", function()
-    it('should return an empty table when no diagnostics are present', function()
-      eq({}, exec_lua [[return vim.lsp.diagnostic.get_line_diagnostics(diagnostic_bufnr, 1)]])
-    end)
-
-    it('should return all diagnostics when no severity is supplied', function()
-      eq(2, exec_lua [[
-        vim.lsp.diagnostic.on_publish_diagnostics(nil, {
-          uri = fake_uri,
-          diagnostics = {
-            make_error("Error 1", 1, 1, 1, 5),
-            make_warning("Warning on Server 1", 1, 1, 2, 5),
-            make_error("Error On Other Line", 2, 1, 1, 5),
-          }
-        }, {client_id=1})
-
-        return #vim.lsp.diagnostic.get_line_diagnostics(diagnostic_bufnr, 1)
-      ]])
-    end)
-
-    it('should return only requested diagnostics when severity_limit is supplied', function()
-      eq(2, exec_lua [[
-        vim.lsp.diagnostic.on_publish_diagnostics(nil, {
-          uri = fake_uri,
-          diagnostics = {
-            make_error("Error 1", 1, 1, 1, 5),
-            make_warning("Warning on Server 1", 1, 1, 2, 5),
-            make_information("Ignored information", 1, 1, 2, 5),
-            make_error("Error On Other Line", 2, 1, 1, 5),
-          }
-        }, {client_id=1})
-
-        return #vim.lsp.diagnostic.get_line_diagnostics(diagnostic_bufnr, 1, { severity_limit = "Warning" })
-      ]])
-    end)
-  end)
-
   describe("vim.lsp.diagnostic.on_publish_diagnostics", function()
-    it('can use functions for config values', function()
-      exec_lua [[
-        vim.lsp.with(vim.lsp.diagnostic.on_publish_diagnostics, {
-          virtual_text = function() return true end,
-        })(nil, {
-            uri = fake_uri,
-            diagnostics = {
-              make_error('Delayed Diagnostic', 4, 4, 4, 4),
-            }
-          }, {client_id=1}
-        )
-      ]]
-
-      eq(1, exec_lua [[return vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1)]])
-      eq(2, exec_lua [[return count_of_extmarks_for_client(diagnostic_bufnr, 1)]])
-
-      -- Now, don't enable virtual text.
-      -- We should have one less extmark displayed.
-      exec_lua [[
-        vim.lsp.with(vim.lsp.diagnostic.on_publish_diagnostics, {
-          virtual_text = function() return false end,
-        })(nil, {
-            uri = fake_uri,
-            diagnostics = {
-              make_error('Delayed Diagnostic', 4, 4, 4, 4),
-            }
-          }, {client_id=1}
-        )
-      ]]
-
-      eq(1, exec_lua [[return vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1)]])
-      eq(1, exec_lua [[return count_of_extmarks_for_client(diagnostic_bufnr, 1)]])
-    end)
-
-    it('can perform updates after insert_leave', function()
-      exec_lua [[vim.api.nvim_set_current_buf(diagnostic_bufnr)]]
-      nvim("input", "o")
-      eq({mode='i', blocking=false}, nvim("get_mode"))
-
-      -- Save the diagnostics
-      exec_lua [[
-        vim.lsp.with(vim.lsp.diagnostic.on_publish_diagnostics, {
-          update_in_insert = false,
-        })(nil, {
-            uri = fake_uri,
-            diagnostics = {
-              make_error('Delayed Diagnostic', 4, 4, 4, 4),
-            }
-          }, {client_id=1}
-        )
-      ]]
-
-      -- No diagnostics displayed yet.
-      eq({mode='i', blocking=false}, nvim("get_mode"))
-      eq(1, exec_lua [[return vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1)]])
-      eq(0, exec_lua [[return count_of_extmarks_for_client(diagnostic_bufnr, 1)]])
-
-      nvim("input", "<esc>")
-      eq({mode='n', blocking=false}, nvim("get_mode"))
-
-      eq(1, exec_lua [[return vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1)]])
-      eq(2, exec_lua [[return count_of_extmarks_for_client(diagnostic_bufnr, 1)]])
-    end)
-
-    it('does not perform updates when not needed', function()
-      exec_lua [[vim.api.nvim_set_current_buf(diagnostic_bufnr)]]
-      nvim("input", "o")
-      eq({mode='i', blocking=false}, nvim("get_mode"))
-
-      -- Save the diagnostics
-      exec_lua [[
-        PublishDiagnostics = vim.lsp.with(vim.lsp.diagnostic.on_publish_diagnostics, {
-          update_in_insert = false,
-          virtual_text = true,
-        })
-
-        -- Count how many times we call display.
-        SetVirtualTextOriginal = vim.diagnostic._set_virtual_text
-
-        DisplayCount = 0
-        vim.diagnostic._set_virtual_text = function(...)
-          DisplayCount = DisplayCount + 1
-          return SetVirtualTextOriginal(...)
-        end
-
-        PublishDiagnostics(nil, {
-            uri = fake_uri,
-            diagnostics = {
-              make_error('Delayed Diagnostic', 4, 4, 4, 4),
-            }
-          }, {client_id=1}
-        )
-      ]]
-
-      -- No diagnostics displayed yet.
-      eq({mode='i', blocking=false}, nvim("get_mode"))
-      eq(1, exec_lua [[return vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1)]])
-      eq(0, exec_lua [[return count_of_extmarks_for_client(diagnostic_bufnr, 1)]])
-      eq(0, exec_lua [[return DisplayCount]])
-
-      nvim("input", "<esc>")
-      eq({mode='n', blocking=false}, nvim("get_mode"))
-
-      eq(1, exec_lua [[return vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1)]])
-      eq(2, exec_lua [[return count_of_extmarks_for_client(diagnostic_bufnr, 1)]])
-      eq(1, exec_lua [[return DisplayCount]])
-
-      -- Go in and out of insert mode one more time.
-      nvim("input", "o")
-      eq({mode='i', blocking=false}, nvim("get_mode"))
-
-      nvim("input", "<esc>")
-      eq({mode='n', blocking=false}, nvim("get_mode"))
-
-      -- Should not have set the virtual text again.
-      eq(1, exec_lua [[return DisplayCount]])
-    end)
-
-    it('never sets virtual text, in combination with insert leave', function()
-      exec_lua [[vim.api.nvim_set_current_buf(diagnostic_bufnr)]]
-      nvim("input", "o")
-      eq({mode='i', blocking=false}, nvim("get_mode"))
-
-      -- Save the diagnostics
-      exec_lua [[
-        PublishDiagnostics = vim.lsp.with(vim.lsp.diagnostic.on_publish_diagnostics, {
-          update_in_insert = false,
-          virtual_text = false,
-        })
-
-        -- Count how many times we call display.
-        SetVirtualTextOriginal = vim.lsp.diagnostic.set_virtual_text
-
-        DisplayCount = 0
-        vim.lsp.diagnostic.set_virtual_text = function(...)
-          DisplayCount = DisplayCount + 1
-          return SetVirtualTextOriginal(...)
-        end
-
-        PublishDiagnostics(nil, {
-            uri = fake_uri,
-            diagnostics = {
-              make_error('Delayed Diagnostic', 4, 4, 4, 4),
-            }
-          }, {client_id=1}
-        )
-      ]]
-
-      -- No diagnostics displayed yet.
-      eq({mode='i', blocking=false}, nvim("get_mode"))
-      eq(1, exec_lua [[return vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1)]])
-      eq(0, exec_lua [[return count_of_extmarks_for_client(diagnostic_bufnr, 1)]])
-      eq(0, exec_lua [[return DisplayCount]])
-
-      nvim("input", "<esc>")
-      eq({mode='n', blocking=false}, nvim("get_mode"))
-
-      eq(1, exec_lua [[return vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1)]])
-      eq(1, exec_lua [[return count_of_extmarks_for_client(diagnostic_bufnr, 1)]])
-      eq(0, exec_lua [[return DisplayCount]])
-
-      -- Go in and out of insert mode one more time.
-      nvim("input", "o")
-      eq({mode='i', blocking=false}, nvim("get_mode"))
-
-      nvim("input", "<esc>")
-      eq({mode='n', blocking=false}, nvim("get_mode"))
-
-      -- Should not have set the virtual text still.
-      eq(0, exec_lua [[return DisplayCount]])
-    end)
-
-    it('can perform updates while in insert mode, if desired', function()
-      exec_lua [[vim.api.nvim_set_current_buf(diagnostic_bufnr)]]
-      nvim("input", "o")
-      eq({mode='i', blocking=false}, nvim("get_mode"))
-
-      -- Save the diagnostics
-      exec_lua [[
-        vim.lsp.with(vim.lsp.diagnostic.on_publish_diagnostics, {
-          update_in_insert = true,
-        })(nil, {
-            uri = fake_uri,
-            diagnostics = {
-              make_error('Delayed Diagnostic', 4, 4, 4, 4),
-            }
-          }, {client_id=1}
-        )
-      ]]
-
-      -- Diagnostics are displayed, because the user wanted them that way!
-      eq({mode='i', blocking=false}, nvim("get_mode"))
-      eq(1, exec_lua [[return vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1)]])
-      eq(2, exec_lua [[return count_of_extmarks_for_client(diagnostic_bufnr, 1)]])
-
-      nvim("input", "<esc>")
-      eq({mode='n', blocking=false}, nvim("get_mode"))
-
-      eq(1, exec_lua [[return vim.lsp.diagnostic.get_count(diagnostic_bufnr, "Error", 1)]])
-      eq(2, exec_lua [[return count_of_extmarks_for_client(diagnostic_bufnr, 1)]])
-    end)
-
     it('allows configuring the virtual text via vim.lsp.with', function()
       local expected_spacing = 10
       local extmarks = exec_lua([[
@@ -710,16 +131,10 @@ describe('vim.lsp.diagnostic', function()
             diagnostics = {
               make_error('Delayed Diagnostic', 4, 4, 4, 4),
             }
-          }, {client_id=1}
+          }, {client_id=client_id}
         )
 
-        return vim.api.nvim_buf_get_extmarks(
-          diagnostic_bufnr,
-          vim.lsp.diagnostic.get_namespace(1),
-          0,
-          -1,
-          { details = true }
-        )
+        return get_extmarks(diagnostic_bufnr, client_id)
       ]], expected_spacing)
 
       local virt_text = extmarks[1][4].virt_text
@@ -728,7 +143,6 @@ describe('vim.lsp.diagnostic', function()
       eq(expected_spacing, #spacing)
     end)
 
-
     it('allows configuring the virtual text via vim.lsp.with using a function', function()
       local expected_spacing = 10
       local extmarks = exec_lua([[
@@ -747,16 +161,10 @@ describe('vim.lsp.diagnostic', function()
             diagnostics = {
               make_error('Delayed Diagnostic', 4, 4, 4, 4),
             }
-          }, {client_id=1}
+          }, {client_id=client_id}
         )
 
-        return vim.api.nvim_buf_get_extmarks(
-          diagnostic_bufnr,
-          vim.lsp.diagnostic.get_namespace(1),
-          0,
-          -1,
-          { details = true }
-        )
+        return get_extmarks(diagnostic_bufnr, client_id)
       ]], expected_spacing)
 
       local virt_text = extmarks[1][4].virt_text
@@ -780,10 +188,10 @@ describe('vim.lsp.diagnostic', function()
               diagnostics = {
                 make_warning('Delayed Diagnostic', 4, 4, 4, 4),
               }
-            }, {client_id=1}
+            }, {client_id=client_id}
           )
 
-          return count_of_extmarks_for_client(diagnostic_bufnr, 1)
+          return #get_extmarks(diagnostic_bufnr, client_id)
         ]], severity_limit)
       end
 
@@ -799,16 +207,6 @@ describe('vim.lsp.diagnostic', function()
       local line = "All 💼 and no 🎉 makes Jack a dull 👦"
       local result = exec_lua([[
         local line = ...
-        local client_id = vim.lsp.start_client {
-          cmd_env = {
-            NVIM_LUA_NOTRACK = "1";
-          };
-          cmd = {
-            vim.v.progpath, '-es', '-u', 'NONE', '--headless'
-          };
-          offset_encoding = "utf-16";
-        }
-
         vim.api.nvim_buf_set_lines(diagnostic_bufnr, 0, -1, false, {line})
 
         vim.lsp.diagnostic.on_publish_diagnostics(nil, {
@@ -829,145 +227,4 @@ describe('vim.lsp.diagnostic', function()
       eq(exec_lua([[return vim.str_byteindex(..., 8, true)]], line), result[1].end_col)
     end)
   end)
-
-  describe('lsp.util.show_line_diagnostics', function()
-    it('creates floating window and returns popup bufnr and winnr if current line contains diagnostics', function()
-      -- Two lines:
-      --    Diagnostic:
-      --    1. <msg>
-      eq(2, exec_lua [[
-        local buffer = vim.api.nvim_create_buf(false, true)
-        vim.api.nvim_buf_set_lines(buffer, 0, -1, false, {
-          "testing";
-          "123";
-        })
-        local diagnostics = {
-          {
-            range = {
-              start = { line = 0; character = 1; };
-              ["end"] = { line = 0; character = 3; };
-            };
-            severity = vim.lsp.protocol.DiagnosticSeverity.Error;
-            message = "Syntax error";
-          },
-        }
-        vim.api.nvim_win_set_buf(0, buffer)
-        vim.lsp.diagnostic.save(diagnostics, buffer, 1)
-        local popup_bufnr, winnr = vim.lsp.diagnostic.show_line_diagnostics()
-        return #vim.api.nvim_buf_get_lines(popup_bufnr, 0, -1, false)
-      ]])
-    end)
-
-    it('creates floating window and returns popup bufnr and winnr without header, if requested', function()
-      -- One line (since no header):
-      --    1. <msg>
-      eq(1, exec_lua [[
-        local buffer = vim.api.nvim_create_buf(false, true)
-        vim.api.nvim_buf_set_lines(buffer, 0, -1, false, {
-          "testing";
-          "123";
-        })
-        local diagnostics = {
-          {
-            range = {
-              start = { line = 0; character = 1; };
-              ["end"] = { line = 0; character = 3; };
-            };
-            severity = vim.lsp.protocol.DiagnosticSeverity.Error;
-            message = "Syntax error";
-          },
-        }
-        vim.api.nvim_win_set_buf(0, buffer)
-        vim.lsp.diagnostic.save(diagnostics, buffer, 1)
-        local popup_bufnr, winnr = vim.lsp.diagnostic.show_line_diagnostics { show_header = false }
-        return #vim.api.nvim_buf_get_lines(popup_bufnr, 0, -1, false)
-      ]])
-    end)
-  end)
-
-  describe('set_signs', function()
-    -- TODO(tjdevries): Find out why signs are not displayed when set from Lua...??
-    pending('sets signs by default', function()
-      exec_lua [[
-        PublishDiagnostics = vim.lsp.with(vim.lsp.diagnostic.on_publish_diagnostics, {
-          update_in_insert = true,
-          signs = true,
-        })
-
-        local diagnostics = {
-          make_error('Delayed Diagnostic', 1, 1, 1, 2),
-          make_error('Delayed Diagnostic', 3, 3, 3, 3),
-        }
-
-        vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
-        vim.lsp.diagnostic.on_publish_diagnostics(nil, {
-            uri = fake_uri,
-            diagnostics = diagnostics
-          }, {client_id=1}
-        )
-
-        vim.lsp.diagnostic.set_signs(diagnostics, diagnostic_bufnr, 1)
-        -- return vim.fn.sign_getplaced()
-      ]]
-
-      nvim("input", "o")
-      nvim("input", "<esc>")
-
-      -- TODO(tjdevries): Find a way to get the signs to display in the test...
-      eq(nil, exec_lua [[
-        return im.fn.sign_getplaced()[1].signs
-      ]])
-    end)
-  end)
-
-  describe('set_loclist()', function()
-    it('sets diagnostics in lnum order', function()
-      local loc_list = exec_lua [[
-        vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
-
-        vim.lsp.diagnostic.on_publish_diagnostics(nil, {
-            uri = fake_uri,
-            diagnostics = {
-              make_error('Farther Diagnostic', 4, 4, 4, 4),
-              make_error('Lower Diagnostic', 1, 1, 1, 1),
-            }
-          }, {client_id=1}
-        )
-
-        vim.lsp.diagnostic.set_loclist()
-
-        return vim.fn.getloclist(0)
-      ]]
-
-      assert(loc_list[1].lnum < loc_list[2].lnum)
-    end)
-
-    it('sets diagnostics in lnum order, regardless of client', function()
-      local loc_list = exec_lua [[
-        vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
-
-        vim.lsp.diagnostic.on_publish_diagnostics(nil, {
-            uri = fake_uri,
-            diagnostics = {
-              make_error('Lower Diagnostic', 1, 1, 1, 1),
-            }
-          }, {client_id=1}
-        )
-
-        vim.lsp.diagnostic.on_publish_diagnostics(nil, {
-            uri = fake_uri,
-            diagnostics = {
-              make_warning('Farther Diagnostic', 4, 4, 4, 4),
-            }
-          }, {client_id=2}
-        )
-
-        vim.lsp.diagnostic.set_loclist()
-
-        return vim.fn.getloclist(0)
-      ]]
-
-      assert(loc_list[1].lnum < loc_list[2].lnum)
-    end)
-  end)
 end)