1 files changed, 50 insertions, 10 deletions
diff --git a/test/functional/lua/diagnostic_spec.lua b/test/functional/lua/diagnostic_spec.lua
index 6b4ca5ac73..17b2d7da4f 100644
--- a/test/functional/lua/diagnostic_spec.lua
+++ b/test/functional/lua/diagnostic_spec.lua
@@ -1042,7 +1042,7 @@ end)
           local virt_text = get_virt_text_extmarks(diagnostic_ns)[1][4].virt_text
 
           local virt_texts = {}
-          for i = 2, #virt_text do
+          for i = 2, #virt_text - 1 do
             table.insert(virt_texts, (string.gsub(virt_text[i][2], "DiagnosticVirtualText", "")))
           end
 
@@ -1086,7 +1086,7 @@ end)
         })
 
         local extmarks = get_virt_text_extmarks(diagnostic_ns)
-        local virt_text = extmarks[1][4].virt_text[2][1]
+        local virt_text = extmarks[1][4].virt_text[3][1]
         return virt_text
       ]]
       eq(' source x: Some error', result)
@@ -1101,7 +1101,7 @@ end)
         }, diagnostic_ns)
 
         local extmarks = get_virt_text_extmarks(diagnostic_ns)
-        local virt_text = extmarks[1][4].virt_text[2][1]
+        local virt_text = extmarks[1][4].virt_text[3][1]
         return virt_text
       ]]
       eq(' Some error', result)
@@ -1121,7 +1121,7 @@ end)
         })
 
         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]}
+        local virt_text = {extmarks[1][4].virt_text[3][1], extmarks[2][4].virt_text[3][1]}
         return virt_text
       ]]
       eq(' source x: Some error', result[1])
@@ -1151,8 +1151,8 @@ end)
         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])
-      eq(" 🔥 Error", result[2][2][1])
+      eq(" 👀 Warning", result[1][3][1])
+      eq(" 🔥 Error", result[2][3][1])
     end)
 
     it('includes source for formatted diagnostics', function()
@@ -1179,8 +1179,48 @@ end)
         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])
-      eq(" another_linter: 🔥 Error", result[2][2][1])
+      eq(" some_linter: 👀 Warning", result[1][3][1])
+      eq(" another_linter: 🔥 Error", result[2][3][1])
+    end)
+
+    it('can add a prefix to virtual text', function()
+      eq('E Some error',  exec_lua [[
+        local diagnostics = {
+          make_error('Some error', 0, 0, 0, 0),
+        }
+
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, diagnostics, {
+          underline = false,
+          virtual_text = {
+            prefix = 'E',
+            suffix = '',
+          }
+        })
+
+        local extmarks = get_virt_text_extmarks(diagnostic_ns)
+        local prefix = extmarks[1][4].virt_text[2][1]
+        local message = extmarks[1][4].virt_text[3][1]
+        return prefix .. message
+      ]])
+
+      eq('[err-code] Some error',  exec_lua [[
+        local diagnostics = {
+          make_error('Some error', 0, 0, 0, 0, nil, 'err-code'),
+        }
+
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, diagnostics, {
+          underline = false,
+          virtual_text = {
+            prefix = function(diag) return string.format('[%s]', diag.code) end,
+            suffix = '',
+          }
+        })
+
+        local extmarks = get_virt_text_extmarks(diagnostic_ns)
+        local prefix = extmarks[1][4].virt_text[2][1]
+        local message = extmarks[1][4].virt_text[3][1]
+        return prefix .. message
+      ]])
     end)
 
     it('can add a suffix to virtual text', function()
@@ -1198,7 +1238,7 @@ end)
         })
 
         local extmarks = get_virt_text_extmarks(diagnostic_ns)
-        local virt_text = extmarks[1][4].virt_text[2][1]
+        local virt_text = extmarks[1][4].virt_text[3][1]
         return virt_text
       ]])
 
@@ -1216,7 +1256,7 @@ end)
         })
 
         local extmarks = get_virt_text_extmarks(diagnostic_ns)
-        local virt_text = extmarks[1][4].virt_text[2][1]
+        local virt_text = extmarks[1][4].virt_text[3][1]
         return virt_text
       ]])
     end)