diff --git a/test/functional/treesitter/highlight_spec.lua b/test/functional/treesitter/highlight_spec.lua
index f2563d9bad..028c01f14a 100644
--- a/test/functional/treesitter/highlight_spec.lua
+++ b/test/functional/treesitter/highlight_spec.lua
@@ -404,7 +404,7 @@ describe('treesitter highlighting (C)', function()
       ]],
     })
 
-    feed 'u'
+    feed 'u:<esc>'
 
     screen:expect({
       grid = [[
@@ -425,7 +425,7 @@ describe('treesitter highlighting (C)', function()
             }                                                            |
           }                                                              |
         ^}                                                                |
-        19 changes; before #2  0 seconds ago                             |
+                                                                         |
       ]],
     })
   end)
diff --git a/test/functional/ui/screen.lua b/test/functional/ui/screen.lua
index 8e15e6c35f..f5cb914299 100644
--- a/test/functional/ui/screen.lua
+++ b/test/functional/ui/screen.lua
@@ -37,10 +37,10 @@
 -- Tests will often share a group of extra attribute sets to expect(). Those can be
 -- defined at the beginning of a test:
 --
---    screen:add_extra_attr_ids {
+--    screen:add_extra_attr_ids({
 --      [100] = { background = Screen.colors.Plum1, underline = true },
 --      [101] = { background = Screen.colors.Red1, bold = true, underline = true },
---    }
+--    })
 --
 -- To help write screen tests, see Screen:snapshot_util().
 -- To debug screen tests, see Screen:redraw_debug().
@@ -454,7 +454,7 @@ end
 ---    screen:expect(grid, [attr_ids])
 ---    screen:expect(condition)
 --- or keyword args (supports more options):
----    screen:expect{grid=[[...]], cmdline={...}, condition=function() ... end}
+---    screen:expect({ grid=[[...]], cmdline={...}, condition=function() ... end })
 ---
 --- @param expected string|function|test.function.ui.screen.Expect
 --- @param attr_ids? table<integer,table<string,any>>
@@ -1713,21 +1713,24 @@ function Screen:_print_snapshot()
       end
     end
     local fn_name = modify_attrs and 'add_extra_attr_ids' or 'set_default_attr_ids'
-    attrstr = ('screen:' .. fn_name .. ' {\n' .. table.concat(attrstrs, '\n') .. '\n}\n\n')
+    attrstr = ('screen:' .. fn_name .. '({\n' .. table.concat(attrstrs, '\n') .. '\n})\n\n')
   end
 
-  local result = ('%sscreen:expect({\n  grid = [[\n  %s\n  ]]'):format(
-    attrstr,
-    kwargs.grid:gsub('\n', '\n  ')
-  )
+  local extstr = ''
   for _, k in ipairs(ext_keys) do
     if ext_state[k] ~= nil and not (k == 'win_viewport' and not self.options.ext_multigrid) then
-      result = result .. ', ' .. k .. '=' .. fmt_ext_state(k, ext_state[k])
+      extstr = extstr .. '\n  ' .. k .. ' = ' .. fmt_ext_state(k, ext_state[k]) .. ','
     end
   end
-  result = result .. '\n})'
 
-  return result
+  return ('%sscreen:expect(%s%s%s%s%s'):format(
+    attrstr,
+    #extstr > 0 and '{\n  grid = [[\n  ' or '[[\n',
+    #extstr > 0 and kwargs.grid:gsub('\n', '\n  ') or kwargs.grid,
+    #extstr > 0 and '\n  ]],' or '\n]]',
+    extstr,
+    #extstr > 0 and '\n})' or ')'
+  )
 end
 
 function Screen:print_snapshot()