1 files changed, 21 insertions, 12 deletions
diff --git a/lua/nvim-treesitter/test_utils.lua b/lua/nvim-treesitter/test_utils.lua
index a19c4048..5665f3ab 100644
--- a/lua/nvim-treesitter/test_utils.lua
+++ b/lua/nvim-treesitter/test_utils.lua
@@ -3,20 +3,20 @@ local M = {}
 local assert = require('luassert')
 
 function M.set_buf_indent_opts(opts)
-    local optnames = {'tabstop', 'shiftwidth', 'softtabstop', 'expandtab'}
-    for _, opt in ipairs(optnames) do
-      if opts[opt] ~= nil then
-        vim.bo[opt] = opts[opt]
-      end
+  local optnames = {'tabstop', 'shiftwidth', 'softtabstop', 'expandtab', 'filetype'}
+  for _, opt in ipairs(optnames) do
+    if opts[opt] ~= nil then
+      vim.bo[opt] = opts[opt]
     end
+  end
 end
 
-function M.indent_whole_file(file, opts)
+function M.run_indent_test(file, runner, opts)
   assert.are.same(1, vim.fn.filereadable(file))
 
   -- load reference file
   vim.cmd(string.format('edit %s', file))
-  local reference = vim.api.nvim_buf_get_lines(0, 0, -1, true)
+  local before = vim.api.nvim_buf_get_lines(0, 0, -1, true)
 
   -- FIXME: why this is not set automatically?
   vim.bo.indentexpr = 'nvim_treesitter#indent()'
@@ -24,15 +24,24 @@ function M.indent_whole_file(file, opts)
 
   M.set_buf_indent_opts(opts)
 
-  -- indent the whole file
-  vim.cmd 'silent normal gg=G'
-  local indented =  vim.api.nvim_buf_get_lines(0, 0, -1, true)
+  -- perform the test
+  runner()
+
+  -- get file content after the test
+  local after = vim.api.nvim_buf_get_lines(0, 0, -1, true)
 
   -- clear any changes to avoid 'No write since last change (add ! to override)'
   vim.cmd 'edit!'
 
-  -- compare before and after
-  assert.are.same(reference, indented)
+  return before, after
+end
+
+function M.indent_whole_file(file, opts)
+  local before, after = M.run_indent_test(file, function()
+    vim.cmd 'silent normal gg=G'
+  end, opts)
+
+  assert.are.same(before, after)
 end
 
 return M