1 files changed, 20 insertions, 12 deletions
diff --git a/test/functional/ui/screen.lua b/test/functional/ui/screen.lua
index ff71194dab..2f2cc85dab 100644
--- a/test/functional/ui/screen.lua
+++ b/test/functional/ui/screen.lua
@@ -181,6 +181,7 @@ end
 -- expected:    Expected screen state (string). Each line represents a screen
 --              row. Last character of each row (typically "|") is stripped.
 --              Common indentation is stripped.
+--              Used as `condition` if NOT a string; must be the ONLY arg then.
 -- attr_ids:    Expected text attributes. Screen rows are transformed according
 --              to this table, as follows: each substring S composed of
 --              characters having the same attributes will be substituted by
@@ -191,18 +192,23 @@ end
 -- any:         true: Succeed if `expected` matches ANY screen line(s).
 --              false (default): `expected` must match screen exactly.
 function Screen:expect(expected, attr_ids, attr_ignore, condition, any)
-  -- remove the last line and dedent
-  expected = dedent(expected:gsub('\n[ ]+$', ''))
   local expected_rows = {}
-  for row in expected:gmatch('[^\n]+') do
-    -- the last character should be the screen delimiter
-    row = row:sub(1, #row - 1)
-    table.insert(expected_rows, row)
-  end
-  if not any then
-    assert(self._height == #expected_rows,
-      "Expected screen state's row count(" .. #expected_rows
-      .. ') differs from configured height(' .. self._height .. ') of Screen.')
+  if type(expected) ~= "string" then
+    assert(not (attr_ids or attr_ignore or condition or any))
+    condition = expected
+    expected = nil
+  else
+    -- Remove the last line and dedent.
+    expected = dedent(expected:gsub('\n[ ]+$', ''))
+    for row in expected:gmatch('[^\n]+') do
+      row = row:sub(1, #row - 1) -- Last char must be the screen delimiter.
+      table.insert(expected_rows, row)
+    end
+    if not any then
+      assert(self._height == #expected_rows,
+        "Expected screen state's row count(" .. #expected_rows
+        .. ') differs from configured height(' .. self._height .. ') of Screen.')
+    end
   end
   local ids = attr_ids or self._default_attr_ids
   local ignore = attr_ignore or self._default_attr_ignore
@@ -218,7 +224,9 @@ function Screen:expect(expected, attr_ids, attr_ignore, condition, any)
       actual_rows[i] = self:_row_repr(self._rows[i], ids, ignore)
     end
 
-    if any then
+    if expected == nil then
+      return
+    elseif any then
       -- Search for `expected` anywhere in the screen lines.
       local actual_screen_str = table.concat(actual_rows, '\n')
       if nil == string.find(actual_screen_str, expected) then