wait for screen (#2584)

* wait for screen

* comments and changelog

* wait for screen after keys

* extra wait for animation

* comment

* comment

* docstring

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - App `title` and `sub_title` attributes can be set to any type https://github.com/Textualize/textual/issues/2521
 - Using `Widget.move_child` where the target and the child being moved are the same is now a no-op https://github.com/Textualize/textual/issues/1743
 - Calling `dismiss` on a screen that is not at the top of the stack now raises an exception https://github.com/Textualize/textual/issues/2575
+- `MessagePump.call_after_refresh` and `MessagePump.call_later` will not return `False` if the callback could not be scheduled. https://github.com/Textualize/textual/pull/2584
 
 ### Fixed
 

src/textual/message_pump.py

@@ -349,20 +349,25 @@ class MessagePump(metaclass=_MessagePumpMeta):
         self._timers.add(timer)
         return timer
 
-    def call_after_refresh(self, callback: Callable, *args: Any, **kwargs: Any) -> None:
+    def call_after_refresh(self, callback: Callable, *args: Any, **kwargs: Any) -> bool:
         """Schedule a callback to run after all messages are processed and the screen
         has been refreshed. Positional and keyword arguments are passed to the callable.
 
         Args:
             callback: A callable.
+
+        Returns:
+            `True` if the callback was scheduled, or `False` if the callback could not be
+                scheduled (may occur if the message pump was closed or closing).
+
         """
         # We send the InvokeLater message to ourselves first, to ensure we've cleared
         # out anything already pending in our own queue.
 
         message = messages.InvokeLater(partial(callback, *args, **kwargs))
-        self.post_message(message)
+        return self.post_message(message)
 
-    def call_later(self, callback: Callable, *args: Any, **kwargs: Any) -> None:
+    def call_later(self, callback: Callable, *args: Any, **kwargs: Any) -> bool:
         """Schedule a callback to run after all messages are processed in this object.
         Positional and keywords arguments are passed to the callable.
 
@@ -370,9 +375,14 @@ class MessagePump(metaclass=_MessagePumpMeta):
             callback: Callable to call next.
             *args: Positional arguments to pass to the callable.
             **kwargs: Keyword arguments to pass to the callable.
+
+        Returns:
+            `True` if the callback was scheduled, or `False` if the callback could not be
+                scheduled (may occur if the message pump was closed or closing).
+
         """
         message = events.Callback(callback=partial(callback, *args, **kwargs))
-        self.post_message(message)
+        return self.post_message(message)
 
     def call_next(self, callback: Callable, *args: Any, **kwargs: Any) -> None:
         """Schedule a callback to run immediately after processing the current message.

src/textual/pilot.py

@@ -65,6 +65,7 @@ class Pilot(Generic[ReturnType]):
         """
         if keys:
             await self._app._press_keys(keys)
+            await self._wait_for_screen()
 
     async def click(
         self,
@@ -132,13 +133,49 @@ class Pilot(Generic[ReturnType]):
         app.post_message(MouseMove(**message_arguments))
         await self.pause()
 
+    async def _wait_for_screen(self, timeout: float = 30.0) -> bool:
+        """Wait for the current screen to have processed all pending events.
+
+        Args:
+            timeout: A timeout in seconds to wait.
+
+        Returns:
+            `True` if all events were processed, or `False` if the wait timed out.
+        """
+        children = [self.app, *self.app.screen.walk_children(with_self=True)]
+        count = 0
+        count_zero_event = asyncio.Event()
+
+        def decrement_counter() -> None:
+            """Decrement internal counter, and set an event if it reaches zero."""
+            nonlocal count
+            count -= 1
+            if count == 0:
+                # When count is zero, all messages queued at the start of the method have been processed
+                count_zero_event.set()
+
+        # Increase the count for every successful call_later
+        for child in children:
+            if child.call_later(decrement_counter):
+                count += 1
+
+        if count:
+            # Wait for the count to return to zero, or a timeout
+            try:
+                await asyncio.wait_for(count_zero_event.wait(), timeout=timeout)
+            except asyncio.TimeoutError:
+                return False
+
+        return True
+
     async def pause(self, delay: float | None = None) -> None:
         """Insert a pause.
 
         Args:
             delay: Seconds to pause, or None to wait for cpu idle.
         """
-        # These sleep zeros, are to force asyncio to give up a time-slice,
+        # These sleep zeros, are to force asyncio to give up a time-slice.
+        await self._wait_for_screen()
         if delay is None:
             await wait_for_idle(0)
         else:
@@ -152,7 +189,9 @@ class Pilot(Generic[ReturnType]):
 
     async def wait_for_scheduled_animations(self) -> None:
         """Wait for any current and scheduled animations to complete."""
+        await self._wait_for_screen()
         await self._app.animator.wait_until_complete()
+        await self._wait_for_screen()
         await wait_for_idle()
         self.app.screen._on_timer_update()
 
@@ -162,5 +201,6 @@ class Pilot(Generic[ReturnType]):
         Args:
             result: The app result returned by `run` or `run_async`.
         """
+        await self._wait_for_screen()
         await wait_for_idle()
         self.app.exit(result)