1 files changed, 25 insertions, 8 deletions
diff --git a/test/functional/ui/screen.lua b/test/functional/ui/screen.lua
index 31669f5578..eb059c38ee 100644
--- a/test/functional/ui/screen.lua
+++ b/test/functional/ui/screen.lua
@@ -165,6 +165,7 @@ function Screen.new(width, height)
     showmode = {},
     showcmd = {},
     ruler = {},
+    hl_groups = {},
     _default_attr_ids = nil,
     _default_attr_ignore = nil,
     _mouse_enabled = true,
@@ -322,7 +323,7 @@ function Screen:expect(expected, attr_ids, attr_ignore)
     assert(not (attr_ids ~= nil or attr_ignore ~= nil))
     local is_key = {grid=true, attr_ids=true, attr_ignore=true, condition=true,
                     any=true, mode=true, unchanged=true, intermediate=true,
-                    reset=true, timeout=true, request_cb=true}
+                    reset=true, timeout=true, request_cb=true, hl_groups=true}
     for _, v in ipairs(ext_keys) do
       is_key[v] = true
     end
@@ -418,9 +419,10 @@ screen:redraw_debug() to show all intermediate screen states.  ]])
     -- (e.g. no external cmdline visible). Some extensions require
     -- preprocessing to represent highlights in a reproducible way.
     local extstate = self:_extstate_repr(attr_state)
-    if expected['mode'] ~= nil then
-      extstate['mode'] = self.mode
+    if expected.mode ~= nil then
+      extstate.mode = self.mode
     end
+
     -- Convert assertion errors into invalid screen state descriptions.
     for _, k in ipairs(concat_tables(ext_keys, {'mode'})) do
       -- Empty states are considered the default and need not be mentioned.
@@ -431,6 +433,17 @@ screen:redraw_debug() to show all intermediate screen states.  ]])
         end
       end
     end
+
+    if expected.hl_groups ~= nil then
+      for name, id in pairs(expected.hl_groups) do
+        local expected_hl = attr_state.ids[id]
+        local actual_hl = self._attr_table[self.hl_groups[name]][(self._options.rgb and 1) or 2]
+        local status, res = pcall(eq, expected_hl, actual_hl, "highlight "..name)
+        if not status then
+          return tostring(res)
+        end
+      end
+    end
   end, expected)
 end
 
@@ -836,6 +849,10 @@ function Screen:_handle_hl_attr_define(id, rgb_attrs, cterm_attrs, info)
   self._new_attrs = true
 end
 
+function Screen:_handle_hl_group_set(name, id)
+  self.hl_groups[name] = id
+end
+
 function Screen:get_hl(val)
   if self._options.ext_newgrid then
     return self._attr_table[val][1]
@@ -1411,17 +1428,17 @@ function Screen:_get_attr_id(attr_state, attrs, hl_id)
     end
     return "UNEXPECTED "..self:_pprint_attrs(self._attr_table[hl_id][1])
   else
-    for id, a in pairs(attr_state.ids) do
-      if self:_equal_attrs(a, attrs) then
-         return id
-       end
-    end
     if self:_equal_attrs(attrs, {}) or
         attr_state.ignore == true or
         self:_attr_index(attr_state.ignore, attrs) ~= nil then
       -- ignore this attrs
       return nil
     end
+    for id, a in pairs(attr_state.ids) do
+      if self:_equal_attrs(a, attrs) then
+         return id
+       end
+    end
     if attr_state.mutable then
       table.insert(attr_state.ids, attrs)
       attr_state.modified = true