Add suspender test for signal set in asyncio task

tests/core/suspend.db

@@ -8,6 +8,9 @@ record(calc, "$(device)COUNTER")  {
 record(ao, "$(device)SIGNAL") {
 }
 
+record(ao, "$(device)STATICSIGNAL") {
+}
+
 record(ao, "$(device)RESUMEVAL") {
     field(VAL, "10")
 }

tests/core/test_suspenders.py

@@ -1,3 +1,4 @@
+import asyncio
 import time
 from pathlib import Path
 from typing import Annotated as A
@@ -16,10 +17,7 @@
 from bluesky.utils import Msg
 
 from ophyd_async.core import SignalRW
-from ophyd_async.epics.core import (
-    EpicsDevice,
-    PvSuffix,
-)
+from ophyd_async.epics.core import EpicsDevice, PvSuffix, epics_signal_rw
 from ophyd_async.epics.testing import TestingIOC
 
 HERE = Path(__file__).resolve().parent
@@ -51,31 +49,27 @@ async def suspend_device(ioc):
 
 
 PARAMETRIZE_SUSPENDERS = pytest.mark.parametrize(
-    "klass,sc_args,resume_val,fail_val,wait_time",
+    "klass,sc_args,resume_val,fail_val,deprecated",
     [
-        (SuspendBoolHigh, (), 0, 1, 0.2),
-        (SuspendBoolLow, (), 1, 0, 0.2),
-        (SuspendFloor, (0.5,), 1, 0, 0.2),
-        (SuspendCeil, (0.5,), 0, 1, 0.2),
-        (SuspendWhenOutsideBand, (0.5, 1.5), 1, 0, 0.2),
-        ((SuspendInBand, True), (0.5, 1.5), 1, 0, 0.2),  # renamed to WhenOutsideBand
-        ((SuspendOutBand, True), (0.5, 1.5), 0, 1, 0.2),
+        (SuspendBoolHigh, (), 0, 1, False),
+        (SuspendBoolLow, (), 1, 0, False),
+        (SuspendFloor, (0.5,), 1, 0, False),
+        (SuspendCeil, (0.5,), 0, 1, False),
+        (SuspendWhenOutsideBand, (0.5, 1.5), 1, 0, False),
+        (SuspendInBand, (0.5, 1.5), 1, 0, True),  # renamed to WhenOutsideBand
+        (SuspendOutBand, (0.5, 1.5), 0, 1, True),
     ],
 )
 
 
 @PARAMETRIZE_SUSPENDERS
 def test_suspender_installed_in_plan(
-    klass, sc_args, resume_val, fail_val, wait_time, RE, suspend_device
+    klass, sc_args, deprecated, resume_val, fail_val, RE, suspend_device
 ):
     sleep_time = 0.2
     fail_time = 0.1
     resume_time = 0.5
-
-    try:
-        klass, deprecated = klass
-    except TypeError:
-        deprecated = False
+    wait_time = 0.2
     if deprecated:
         with pytest.warns(UserWarning):
             suspender = klass(
@@ -108,22 +102,12 @@ def plan():
 
 @PARAMETRIZE_SUSPENDERS
 async def test_suspender_installed_outside_plan(
-    klass, sc_args, resume_val, fail_val, wait_time, RE, suspend_device
+    klass, sc_args, deprecated, resume_val, fail_val, RE, suspend_device
 ):
     sleep_time = 0.2
     fail_time = 0.1
     resume_time = 0.5
-
-    await suspend_device.resume_signal.set(resume_val, wait=True)
-    await suspend_device.fail_signal.set(fail_val, wait=True)
-    await suspend_device.resume_time_signal.set(resume_time, wait=True)
-    await suspend_device.fail_time_signal.set(fail_time, wait=True)
-    await suspend_device.signal.set(resume_val, wait=True)
-
-    try:
-        klass, deprecated = klass
-    except TypeError:
-        deprecated = False
+    wait_time = 0.2
     if deprecated:
         with pytest.warns(UserWarning):
             suspender = klass(
@@ -136,9 +120,63 @@ async def test_suspender_installed_outside_plan(
 
     RE.install_suspender(suspender)
     scan = [Msg("checkpoint"), Msg("sleep", None, sleep_time)]
+    await suspend_device.resume_signal.set(resume_val, wait=True)
+    await suspend_device.fail_signal.set(fail_val, wait=True)
+    await suspend_device.resume_time_signal.set(resume_time, wait=True)
+    await suspend_device.fail_time_signal.set(fail_time, wait=True)
+    await suspend_device.signal.set(resume_val, wait=True)
     await suspend_device.counter_signal.set(0, wait=True)
     start = time.time()
     RE(scan)
     stop = time.time()
     delta = stop - start
     assert delta >= resume_time + sleep_time + wait_time
+
+
+@PARAMETRIZE_SUSPENDERS
+@pytest.mark.parametrize("task_in_plan", [True, False])
+async def test_suspension_from_async_task(
+    klass, sc_args, deprecated, resume_val, fail_val, task_in_plan, RE, ioc
+):
+    sleep_time = 0.2
+    fail_time = 0.1
+    resume_time = 0.5
+    wait_time = 0.2
+
+    signal = epics_signal_rw(
+        float, f"{ioc.prefix_for(SuspenderTestDevice)}STATICSIGNAL"
+    )
+    await signal.connect()
+    await signal.set(resume_val)  # set to initial non-suspending value
+
+    if deprecated:
+        with pytest.warns(UserWarning):
+            suspender = klass(signal, *sc_args, sleep=wait_time, is_async=True)
+    else:
+        suspender = klass(signal, *sc_args, sleep=wait_time, is_async=True)
+
+    async def _set_after_time(delay, value):
+        await asyncio.sleep(delay)
+        await signal.set(value)
+
+    tasks = []
+
+    RE.install_suspender(suspender)
+    if not task_in_plan:
+        tasks.append(asyncio.create_task(_set_after_time(fail_time, fail_val)))
+        tasks.append(asyncio.create_task(_set_after_time(resume_time, resume_val)))
+
+    def _plan():
+        if task_in_plan:
+            tasks.append(asyncio.create_task(_set_after_time(fail_time, fail_val)))
+            tasks.append(asyncio.create_task(_set_after_time(resume_time, resume_val)))
+        yield from bps.checkpoint()
+        yield from bps.sleep(sleep_time)
+
+    start = time.time()
+    RE(_plan())
+    stop = time.time()
+    for task in tasks:
+        await task
+    delta = stop - start
+    assert delta >= resume_time + sleep_time + wait_time