tests

doc/plan_stubs/matplotlib_helpers.md

@@ -36,10 +36,10 @@ from bluesky.preprocessors import subs_decorator
 def my_plan():
     # BAD
     # fig, ax = plt.subplots()
-    
+
     # GOOD
     fig, ax = yield from matplotlib_subplots()
-    
+
     # Pass the matplotlib ax object to other callbacks
     @subs_decorator([
         LiveFitPlot(..., ax=ax),

src/ibex_bluesky_core/plan_stubs/__init__.py

@@ -51,7 +51,8 @@ def call_sync(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> Genera
 
 
 def matplotlib_subplots(
-    *args: list[Any], **kwargs: dict[Any]
+    *args: Any,  # noqa: ANN401 - pyright doesn't understand we're wrapping mpl API.
+    **kwargs: Any,  # noqa: ANN401 - pyright doesn't understand we're wrapping mpl API.
 ) -> Generator[Msg, None, tuple[Figure, Any]]:
     """Create a new matplotlib figure and axes, using matplotlib.pyplot.subplots, from a plan.
 

src/ibex_bluesky_core/run_engine/_msg_handlers.py

@@ -109,8 +109,8 @@ async def call_qt_safe_handler(msg: Msg) -> Any:  # noqa: ANN401
     we would have no way to prove that for a generic user-supplied function. So we only
     expose known cases, like matplotlib.pyplot.subplots for example.
 
-    In particular, operations that take longer than CALL_QT_SAFE_TIMEOUT will cause an
-    asyncio TimeoutError (to flag obvious misuse), but won't cancel the underlying task.
+    Slow functions will cause various undesirable side effects, like stalling the event loop
+    and interfering with ctrl-c handling.
 
     """
     func = msg.obj
@@ -140,7 +140,13 @@ def start(self, doc: RunStart) -> None:
     # Send fake event to our callback to trigger it (actual contents unimportant)
     cb("start", {"time": 0, "uid": ""})
 
-    await wait_for(done_event.wait(), CALL_QT_SAFE_TIMEOUT)
+    try:
+        await wait_for(done_event.wait(), CALL_QT_SAFE_TIMEOUT)
+    except TimeoutError as e:
+        raise TimeoutError(
+            f"Long-running function '{func.__name__}' passed to call_qt_safe_handler. Functions "
+            f"passed to call_qt_safe_handler must be faster than {CALL_QT_SAFE_TIMEOUT}s."
+        ) from e
     if exc is not None:
         raise exc
     return result

tests/test_plan_stubs.py

@@ -1,13 +1,13 @@
 # pyright: reportMissingParameterType=false
-
+import re
 import time
 from asyncio import CancelledError
 from unittest.mock import patch
 
 import pytest
 from bluesky.utils import Msg
 
-from ibex_bluesky_core.plan_stubs import call_sync
+from ibex_bluesky_core.plan_stubs import CALL_QT_SAFE_MSG_KEY, call_sync, matplotlib_subplots
 from ibex_bluesky_core.run_engine._msg_handlers import call_sync_handler
 
 
@@ -66,3 +66,55 @@ def f():
     end = time.monotonic()
 
     assert end - start == pytest.approx(1, abs=0.2)
+
+
+def test_call_qt_safe_returns_result(RE):
+    def f(arg, keyword_arg):
+        assert arg == "foo"
+        assert keyword_arg == "bar"
+        return 123
+
+    def plan():
+        return (yield Msg(CALL_QT_SAFE_MSG_KEY, f, "foo", keyword_arg="bar"))
+
+    result = RE(plan())
+
+    assert result.plan_result == 123
+
+
+def test_call_qt_safe_throws_exception(RE):
+    def f():
+        raise ValueError("broke it")
+
+    def plan():
+        return (yield Msg(CALL_QT_SAFE_MSG_KEY, f))
+
+    with pytest.raises(ValueError, match="broke it"):
+        RE(plan())
+
+
+def test_call_qt_safe_blocking_causes_descriptive_timeout_error(RE):
+    def f():
+        time.sleep(0.2)
+
+    def plan():
+        return (yield Msg(CALL_QT_SAFE_MSG_KEY, f))
+
+    with patch("ibex_bluesky_core.run_engine._msg_handlers.CALL_QT_SAFE_TIMEOUT", new=0.1):
+        with pytest.raises(
+            TimeoutError,
+            match=re.escape(
+                "Long-running function 'f' passed to call_qt_safe_handler. "
+                "Functions passed to call_qt_safe_handler must be faster than 0.1s."
+            ),
+        ):
+            RE(plan())
+
+
+def test_matplotlib_subplots_calls_pyplot_subplots(RE):
+    def plan():
+        return (yield from matplotlib_subplots("foo", keyword="bar"))
+
+    with patch("ibex_bluesky_core.plan_stubs.plt.subplots") as mock_pyplot_subplots:
+        RE(plan())
+        mock_pyplot_subplots.assert_called_once_with("foo", keyword="bar")