1 files changed, 52 insertions, 6 deletions
diff --git a/lua/nvim-treesitter/test_utils.lua b/lua/nvim-treesitter/test_utils.lua
index 9059aef4e..264f5b1db 100644
--- a/lua/nvim-treesitter/test_utils.lua
+++ b/lua/nvim-treesitter/test_utils.lua
@@ -1,6 +1,55 @@
 local M = {}
 
 local assert = require('luassert')
+local say = require('say')
+
+local function same_indent(state, arguments)
+  local before = arguments[1]
+  local after = arguments[2]
+
+  local ok = true
+  errors = { before = {}, after = {} }
+  for line = 1,#before do
+    if before[line] ~= after[line] then
+      errors.before[line] = #string.match(before[line], '^%s*')
+      errors.after[line] = #string.match(after[line], '^%s*')
+      ok = false
+    end
+  end
+
+  -- swap arguments to be consistent with all the assertions in plenary.luassert
+  table.insert(arguments, 1, table.remove(arguments, 2))
+  arguments.fmtargs = {}
+  arguments.fmtargs[1] = { errors = errors.after }
+  arguments.fmtargs[2] = { errors = errors.before }
+
+  return ok
+end
+
+local function format_indent(arg, fmtargs)
+  local output = {}
+  for i, line in ipairs(arg) do
+    if fmtargs.errors[i] then
+      table.insert(output, string.format('%2d => |%s', fmtargs.errors[i], line))
+    else
+      table.insert(output, string.format('      |%s', line))
+    end
+  end
+  return table.concat(output, '\n')
+end
+
+say:set_namespace('en')
+say:set('assertion.same_indent.positive', 'Expected indentation to be the same.\nFound:\n%s\nExpected:\n%s')
+say:set('assertion.same_indent.negative', 'Expected indentation to be different.\nFound:\n%s\nExpected:\n%s')
+assert:register('assertion', 'same_indent', same_indent,
+  'assertion.same_indent.positive', 'assert.same_indent.negative')
+
+-- Custom assertion better suited for indentation diffs
+local function compare_indent(before, after)
+  assert:add_formatter(format_indent)
+  assert.is.same_indent(before, after)
+  assert:remove_formatter(format_indent)
+end
 
 function M.set_buf_indent_opts(opts)
   local optnames = {'tabstop', 'shiftwidth', 'softtabstop', 'expandtab', 'filetype'}
@@ -41,7 +90,7 @@ function M.indent_whole_file(file, opts)
     vim.cmd 'silent normal gg=G'
   end, opts)
 
-  assert.are.same(before, after)
+  compare_indent(before, after)
 end
 
 -- Open a file, use `normal o` to insert a new line and compare results
@@ -59,12 +108,9 @@ function M.indent_new_line(file, spec, opts)
   end, opts)
 
   local indent = type(spec.indent) == 'string' and spec.indent or string.rep(' ', spec.indent)
-
-  -- print('# before:\n', table.concat(before, '\n'))
-  -- print('# after:\n', table.concat(after, '\n'))
-
   table.insert(before, spec.on_line + 1, indent .. spec.text)
-  assert.are.same(before, after)
+
+  compare_indent(before, after)
 end
 
 return M