1 files changed, 24 insertions, 23 deletions
diff --git a/tests/query/highlights_spec.lua b/tests/query/highlights_spec.lua
index 3bb03d49d..b6a52124a 100644
--- a/tests/query/highlights_spec.lua
+++ b/tests/query/highlights_spec.lua
@@ -1,39 +1,39 @@
-local highlighter = require "vim.treesitter.highlighter"
-local parsers = require "nvim-treesitter.parsers"
+local highlighter = require('vim.treesitter.highlighter')
 local ts = vim.treesitter
 
 local COMMENT_NODES = {
-  markdown = "html_block",
-  haskell = "haddock",
+  markdown = 'html_block',
+  haskell = 'haddock',
 }
 
 local function check_assertions(file)
   local buf = vim.fn.bufadd(file)
   vim.fn.bufload(file)
-  local lang = parsers.get_buf_lang(buf)
+  local ft = vim.bo[buf].filetype
+  local lang = vim.treesitter.language.get_lang(ft) or ft
   assert.same(
     1,
-    vim.fn.executable "highlight-assertions",
+    vim.fn.executable('highlight-assertions'),
     '"highlight-assertions" not executable!'
       .. ' Get it via "cargo install --git https://github.com/theHamsta/highlight-assertions"'
   )
-  local comment_node = COMMENT_NODES[lang] or "comment"
+  local comment_node = COMMENT_NODES[lang] or 'comment'
   local assertions = vim.fn.json_decode(
     vim.fn.system(
       "highlight-assertions -p '"
-        .. vim.api.nvim_get_runtime_file("parser/" .. lang .. ".so", false)[1]
+        .. vim.api.nvim_get_runtime_file('parser/' .. lang .. '.so', false)[1]
         .. "' -s '"
         .. file
         .. "' -c "
         .. comment_node
     )
   )
-  local parser = parsers.get_parser(buf, lang)
+  local parser = ts.get_parser(buf, lang)
   parser:parse(true)
 
   local self = highlighter.new(parser, {})
 
-  assert.True(#assertions > 0, "No assertions detected!")
+  assert.True(#assertions > 0, 'No assertions detected!')
   for _, assertion in ipairs(assertions) do
     local row = assertion.position.row
     local col = assertion.position.column
@@ -46,7 +46,7 @@ local function check_assertions(file)
         return
       end
 
-      local root = state.tstree:root()
+      local root = tstree:root()
       local root_start_row, _, root_end_row, _ = root:range()
 
       -- Only worry about trees within the line range
@@ -72,21 +72,22 @@ local function check_assertions(file)
         assert.is.number(col)
         if hl and ts.is_in_node_range(node, row, col) then
           local c = query._query.captures[capture] -- name of the capture in the query
-          if c ~= nil and c ~= "spell" and c ~= "conceal" then
+          if c ~= nil and c ~= 'spell' and c ~= 'conceal' then
             captures[c] = true
             highlights[c] = true
           end
         end
       end
-    end, true)
-    if assertion.expected_capture_name:match "^!" then
+    end)
+    if assertion.expected_capture_name:match('^!') then
       assert.Falsy(
-        captures[assertion.expected_capture_name:sub(2)] or highlights[assertion.expected_capture_name:sub(2)],
-        "Error in at "
+        captures[assertion.expected_capture_name:sub(2)]
+          or highlights[assertion.expected_capture_name:sub(2)],
+        'Error in at '
           .. file
-          .. ":"
+          .. ':'
           .. (row + 1)
-          .. ":"
+          .. ':'
           .. (col + 1)
           .. ': expected "'
           .. assertion.expected_capture_name
@@ -98,11 +99,11 @@ local function check_assertions(file)
     else
       assert.True(
         captures[assertion.expected_capture_name] or highlights[assertion.expected_capture_name],
-        "Error in at "
+        'Error in at '
           .. file
-          .. ":"
+          .. ':'
           .. (row + 1)
-          .. ":"
+          .. ':'
           .. (col + 1)
           .. ': expected "'
           .. assertion.expected_capture_name
@@ -115,8 +116,8 @@ local function check_assertions(file)
   end
 end
 
-describe("highlight queries", function()
-  local files = vim.fn.split(vim.fn.glob "tests/query/highlights/**/*.*")
+describe('highlight queries', function()
+  local files = vim.fn.split(vim.fn.glob('tests/query/highlights/**/*.*'))
   for _, file in ipairs(files) do
     it(file, function()
       check_assertions(file)