test: fix tests broken after 0.10 update (#6714)

2 files changed, 10 insertions, 9 deletions

diff --git a/tests/query/highlights_spec.lua b/tests/query/highlights_spec.lua
index ff6a7523d..3bb03d49d 100644
--- a/tests/query/highlights_spec.lua
+++ b/tests/query/highlights_spec.lua
@@ -29,6 +29,7 @@ local function check_assertions(file)
     )
   )
   local parser = parsers.get_parser(buf, lang)
+  parser:parse(true)
 
   local self = highlighter.new(parser, {})
 
@@ -39,12 +40,13 @@ local function check_assertions(file)
 
     local captures = {}
     local highlights = {}
-    self.tree:for_each_tree(function(tstree, tree)
-      if not tstree then
+    self:prepare_highlight_states(row, row + 1)
+    self:for_each_highlight_state(function(state)
+      if not state.tstree then
         return
       end
 
-      local root = tstree:root()
+      local root = state.tstree:root()
       local root_start_row, _, root_end_row, _ = root:range()
 
       -- Only worry about trees within the line range
@@ -52,7 +54,7 @@ local function check_assertions(file)
         return
       end
 
-      local query = self:get_query(tree:lang())
+      local query = state.highlighter_query
 
       -- Some injected languages may not have highlight queries.
       if not query:query() then
@@ -62,7 +64,7 @@ local function check_assertions(file)
       local iter = query:query():iter_captures(root, self.bufnr, row, row + 1)
 
       for capture, node, _ in iter do
-        local hl = query.hl_cache[capture]
+        local hl = query:get_hl_from_capture(capture)
         assert.is.truthy(hl)
 
         assert.Truthy(node)
diff --git a/tests/query/injection_spec.lua b/tests/query/injection_spec.lua
index 7efbab315..95ad12eec 100644
--- a/tests/query/injection_spec.lua
+++ b/tests/query/injection_spec.lua
@@ -1,5 +1,4 @@
 require "nvim-treesitter.highlight" -- yes, this is necessary to set the hlmap
-local highlighter = require "vim.treesitter.highlighter"
 local configs = require "nvim-treesitter.configs"
 local parsers = require "nvim-treesitter.parsers"
 local ts = vim.treesitter
@@ -21,8 +20,8 @@ local function check_assertions(file)
   )
   local parser = parsers.get_parser(buf, lang)
 
-  local self = highlighter.new(parser, {})
-  local top_level_root = parser:parse()[1]:root()
+  local self = parser
+  local top_level_root = parser:parse(true)[1]:root()
 
   for _, assertion in ipairs(assertions) do
     local row = assertion.position.row
@@ -32,7 +31,7 @@ local function check_assertions(file)
     assertion.expected_capture_name = neg_assert and assertion.expected_capture_name:sub(2)
       or assertion.expected_capture_name
     local found = false
-    self.tree:for_each_tree(function(tstree, tree)
+    self:for_each_tree(function(tstree, tree)
       if not tstree then
         return
       end