Repository: bluesky/bluesky
Pull request #1758 (2024-06-28 09:28:15)

returned the results of `asyncio.wait` in `bps.wait_for`
Also updated docs to reflect this, and added a test.



Patch:
diff --git a/docs/msg.rst b/docs/msg.rst
index ce09e43c..b9d0db98 100644
--- a/docs/msg.rst
+++ b/docs/msg.rst
@@ -263,6 +263,17 @@ Ex ::
     future.set_result(3)
     future.done() # will give True
 
+Returns a set of the tasks that were waited on.
+Ex ::
+
+    from asyncio.futures import Future
+    future = Future()
+    future.done() # will give false
+    def plan()
+        results = (yield Msg('wait_for', [lambda : future ,]))
+        for result in results:
+            # will give the exception raised
+            result.exception()
 
 input
 +++++
diff --git a/src/bluesky/run_engine.py b/src/bluesky/run_engine.py
index 8d30c2a3..09dca6e7 100644
--- a/src/bluesky/run_engine.py
+++ b/src/bluesky/run_engine.py
@@ -1781,7 +1781,7 @@ class RunEngine:
         return plan_return
 
     async def _wait_for(self, msg):
-        """Instruct the RunEngine to wait for futures
+        """Instruct the RunEngine to wait for futures and return the resulting tasks.
 
         Expected message object is:
 
@@ -1798,9 +1798,10 @@ class RunEngine:
 
         (futs,) = msg.args
         futs = [asyncio.ensure_future(f()) for f in futs]
-        _, pending = await asyncio.wait(futs, **self._loop_for_kwargs, **msg.kwargs)
+        completed, pending = await asyncio.wait(futs, **self._loop_for_kwargs, **msg.kwargs)
         if pending:
             raise WaitForTimeoutError("Plan failed to complete in the specified time")
+        return futs
 
     async def _open_run(self, msg):
         """Instruct the RunEngine to start a new "run"


Test patch:
diff --git a/src/bluesky/tests/test_run_engine.py b/src/bluesky/tests/test_run_engine.py
index 85a21060..0d5ddac2 100644
--- a/src/bluesky/tests/test_run_engine.py
+++ b/src/bluesky/tests/test_run_engine.py
@@ -13,7 +13,7 @@ import pytest
 from event_model import DocumentNames
 
 from bluesky import Msg
-from bluesky.plan_stubs import abs_set, checkpoint, declare_stream, pause, trigger_and_read, wait
+from bluesky.plan_stubs import abs_set, checkpoint, declare_stream, pause, trigger_and_read, wait, wait_for
 from bluesky.plans import count, grid_scan
 from bluesky.preprocessors import (
     SupplementalData,
@@ -35,6 +35,7 @@ from bluesky.run_engine import (
     RunEngineInterrupted,
     RunEngineStateMachine,
     TransitionError,
+    WaitForTimeoutError,
 )
 from bluesky.tests import requires_ophyd
 from bluesky.tests.utils import DocCollector, MsgCollector
@@ -1956,6 +1957,63 @@ def test_wait_with_timeout(set_finished, RE):
             RE(plan())
 
 
+async def passing_coroutine():
+    await asyncio.sleep(0.001)
+    return 1080
+
+
+async def failing_coroutine():
+    await asyncio.sleep(0.001)
+    raise RuntimeError("This is a runtime error")
+
+
+def test_wait_with_planstub_coroutine(RE):
+    def runtime_error_and_success_plan():
+        tasks = yield from wait_for([lambda: passing_coroutine(), lambda: failing_coroutine()])
+        passing_future, failing_future = tasks
+        assert isinstance(passing_future, asyncio.Task) and isinstance(failing_future, asyncio.Task)
+        assert passing_future.get_coro().__name__ == "passing_coroutine"
+        assert failing_future.get_coro().__name__ == "failing_coroutine"
+        assert passing_future.done() and passing_future.result() == 1080
+        assert failing_future.done() and isinstance(failing_future.exception(), RuntimeError)
+
+        with pytest.raises(RuntimeError, match="This is a runtime error"):
+            raise failing_future.exception()  # the error is the same
+
+    RE(runtime_error_and_success_plan())
+
+
+def test_wait_with_planstub_status(RE):
+    def runtime_error_and_success_plan():
+        passing_task = asyncio.create_task(passing_coroutine())
+
+        failing_task = asyncio.create_task(failing_coroutine())
+
+        futures = yield from wait_for([lambda: passing_task, lambda: failing_task])
+        assert passing_task is futures[0] and failing_task is futures[1]
+        assert passing_task.done() and passing_task.result() == 1080
+        assert failing_task.done() and isinstance(failing_task.exception(), RuntimeError)
+
+    RE(runtime_error_and_success_plan())
+
+
+def test_wait_with_planstub_timeout(RE):
+    def timeout_plan():
+        async def timeout_coro():
+            await asyncio.sleep(1e9)
+
+        timeout_task = asyncio.create_task(timeout_coro())
+        try:
+            yield from wait_for([lambda: timeout_task], timeout=0.01)
+        finally:
+            # Manually tear down the tasks, otherwise it can happen
+            # after the test succeeds/fails
+            timeout_task.cancel()
+
+    with pytest.raises(WaitForTimeoutError, match="Plan failed to complete in the specified time"):
+        RE(timeout_plan())
+
+
 def test_1event_rewind(RE, hw):
     cache = []