Don't return completed_status() in set_and_wait_for_value (#805)

* Don't return completed_status() in set_and_wait_for_value

* Add support for callable match_value in set_and_wait_for_value

* explicitly use asyncio.TimeoutError in set_and_wait_for_other_value

Author: jsouter

src/ophyd_async/core/_signal.py

@@ -23,7 +23,7 @@
 from ._protocol import AsyncReadable, AsyncStageable
 from ._signal_backend import SignalBackend, SignalDatatypeT, SignalDatatypeV
 from ._soft_signal_backend import SoftSignalBackend
-from ._status import AsyncStatus, completed_status
+from ._status import AsyncStatus
 from ._utils import (
     CALCULATE_TIMEOUT,
     DEFAULT_TIMEOUT,
@@ -579,25 +579,34 @@ async def set_and_wait_for_other_value(
 
     status = set_signal.set(set_value, timeout=set_timeout)
 
+    if callable(match_value):
+        matcher: Callable[[SignalDatatypeV], bool] = match_value  # type: ignore
+    else:
+
+        def matcher(value):
+            return value == match_value
+
+        matcher.__name__ = f"equals_{match_value}"
+
     # If the value was the same as before no need to wait for it to change
-    if current_value != match_value:
+    if not matcher(current_value):
 
         async def _wait_for_value():
             async for value in values_gen:
-                if value == match_value:
+                if matcher(value):
                     break
 
         try:
             await asyncio.wait_for(_wait_for_value(), timeout)
             if wait_for_set_completion:
                 await status
-            return status
         except asyncio.TimeoutError as e:
-            raise TimeoutError(
-                f"{match_signal.name} didn't match {match_value} in {timeout}s"
+            raise asyncio.TimeoutError(
+                f"{match_signal.name} value didn't match value from"
+                f" {matcher.__name__}() in {timeout}s"
             ) from e
 
-    return completed_status()
+    return status
 
 
 async def set_and_wait_for_value(

tests/core/test_signal.py

@@ -239,10 +239,79 @@ async def check_set_and_wait():
             )
         )
 
-    with pytest.raises(TimeoutError):
+    with pytest.raises(asyncio.TimeoutError):
         await asyncio.gather(wait_and_set_read(), check_set_and_wait())
 
 
+async def test_status_of_set_and_wait_for_value():
+    set_signal = epics_signal_rw(int, "pva://signal")
+    match_signal = epics_signal_rw(int, "pva://match_signal")
+
+    async def set_match_signal_after_delay(value: Any, **kwargs):
+        await asyncio.sleep(0.3)
+        await match_signal.set(value / 2)
+
+    await set_signal.connect(mock=True)
+    await match_signal.connect(mock=True)
+    callback_on_mock_put(set_signal, set_match_signal_after_delay)  # type: ignore
+
+    status = await set_and_wait_for_value(set_signal, 2)
+    assert status.done
+
+    with pytest.raises(asyncio.TimeoutError):
+        status = await set_and_wait_for_value(set_signal, 4, timeout=0.1)
+
+    with pytest.raises(asyncio.TimeoutError):
+        status = await set_and_wait_for_value(
+            set_signal, 8, timeout=0.1, wait_for_set_completion=False
+        )
+
+    status = await set_and_wait_for_other_value(set_signal, 32, match_signal, 16)
+    assert status.done
+
+    status = await set_and_wait_for_other_value(
+        set_signal, 30, match_signal, 15, wait_for_set_completion=False
+    )
+    assert not status.done
+    await status
+
+    with pytest.raises(asyncio.TimeoutError):
+        status = await set_and_wait_for_other_value(
+            set_signal, 30, match_signal, -1, timeout=0.5
+        )
+
+
+async def test_callable_match_value_set_and_wait_for_value():
+    set_signal = epics_signal_rw(int, "pva://signal")
+    match_signal = epics_signal_rw(int, "pva://match_signal")
+
+    async def set_match_signal_after_delay(value: Any, **kwargs):
+        await asyncio.sleep(0.3)
+        await match_signal.set(value / 2)
+
+    await set_signal.connect(mock=True)
+    await match_signal.connect(mock=True)
+    callback_on_mock_put(set_signal, set_match_signal_after_delay)  # type: ignore
+
+    def _equals_x(value, x):
+        return value == x
+
+    status = await set_and_wait_for_value(
+        set_signal, 10, lambda val: _equals_x(val, 10)
+    )
+    assert status.done
+
+    status = await set_and_wait_for_other_value(
+        set_signal, 20, match_signal, lambda val: _equals_x(val, 10)
+    )
+    assert status.done
+
+    with pytest.raises(asyncio.TimeoutError):
+        status = await set_and_wait_for_other_value(
+            set_signal, 30, match_signal, lambda val: _equals_x(val, -1), timeout=0.5
+        )
+
+
 async def test_wait_for_value_with_value():
     signal = epics_signal_rw(str, read_pv="pva://signal", name="signal")
     await signal.connect(mock=True)

tests/epics/eiger/test_eiger_controller.py

@@ -1,3 +1,4 @@
+import asyncio
 from unittest.mock import ANY, patch
 
 from pytest import fixture, raises
@@ -76,7 +77,7 @@ async def test_given_detector_fails_to_go_ready_when_arm_called_then_fails(
     eiger_driver_and_controller_no_arm: DriverAndController,
 ):
     _, controller = eiger_driver_and_controller_no_arm
-    with raises(TimeoutError):
+    with raises(asyncio.TimeoutError):
         await controller.prepare(TriggerInfo(number_of_triggers=10))
         await controller.arm()
         await controller.wait_for_idle()