1 files changed, 31 insertions, 1 deletions
diff --git a/test/functional/ui/screen.lua b/test/functional/ui/screen.lua
index e8a39ab6f8..b3d1abca69 100644
--- a/test/functional/ui/screen.lua
+++ b/test/functional/ui/screen.lua
@@ -179,6 +179,7 @@ function Screen.new(width, height)
     _width = width,
     _height = height,
     _grids = {},
+    _grid_win_extmarks = {},
     _cursor = {
       grid = 1, row = 1, col = 1
     },
@@ -278,6 +279,8 @@ local ext_keys = {
 --              attributes in the final state are an error.
 --              Use screen:set_default_attr_ids() to define attributes for many
 --              expect() calls.
+-- extmarks:    Expected win_extmarks accumulated for the grids. For each grid,
+--              the win_extmark messages are accumulated into an array.
 -- condition:   Function asserting some arbitrary condition. Return value is
 --              ignored, throw an error (use eq() or similar) to signal failure.
 -- any:         Lua pattern string expected to match a screen line. NB: the
@@ -320,7 +323,7 @@ function Screen:expect(expected, attr_ids, ...)
     assert(not (attr_ids ~= nil))
     local is_key = {grid=true, attr_ids=true, condition=true, mouse_enabled=true,
                     any=true, mode=true, unchanged=true, intermediate=true,
-                    reset=true, timeout=true, request_cb=true, hl_groups=true}
+                    reset=true, timeout=true, request_cb=true, hl_groups=true, extmarks=true}
     for _, v in ipairs(ext_keys) do
       is_key[v] = true
     end
@@ -459,6 +462,25 @@ screen:redraw_debug() to show all intermediate screen states.  ]])
         end
       end
     end
+
+    if expected.extmarks ~= nil then
+      for gridid, expected_marks in pairs(expected.extmarks) do
+        local stored_marks = self._grid_win_extmarks[gridid]
+        if stored_marks == nil then
+          return 'no win_extmark for grid '..tostring(gridid)
+        end
+        local status, res = pcall(eq, expected_marks, stored_marks, "extmarks for grid "..tostring(gridid))
+        if not status then
+          return tostring(res)
+        end
+      end
+      for gridid, _ in pairs(self._grid_win_extmarks) do
+        local expected_marks = expected.extmarks[gridid]
+        if expected_marks == nil then
+          return 'unexpected win_extmark for grid '..tostring(gridid)
+        end
+      end
+    end
   end, expected)
 end
 
@@ -703,6 +725,7 @@ function Screen:_reset()
   self.cmdline_block = {}
   self.wildmenu_items = nil
   self.wildmenu_pos = nil
+  self._grid_win_extmarks = {}
 end
 
 function Screen:_handle_mode_info_set(cursor_style_enabled, mode_info)
@@ -803,6 +826,13 @@ function Screen:_handle_win_close(grid)
   self.float_pos[grid] = nil
 end
 
+function Screen:_handle_win_extmark(grid, ...)
+  if self._grid_win_extmarks[grid] == nil then
+    self._grid_win_extmarks[grid] = {}
+  end
+  table.insert(self._grid_win_extmarks[grid], {...})
+end
+
 function Screen:_handle_busy_start()
   self._busy = true
 end