test: fix tests after refactoring

tests/tests_devices/test_devices_falcon.py

@@ -5,6 +5,7 @@ from unittest import mock
 
 import ophyd
 import pytest
+from ophyd.utils.errors import StatusTimeoutError
 from ophyd_devices.tests.utils import MockPV, patch_dual_pvs
 
 from superxas_bec.devices.falcon import FalconAcquiringStatus, FalconSuperXAS
@@ -50,8 +51,9 @@ def test_devices_falcon_stage(falcon):
     # Should timeout
     falcon.acquiring.put(FalconAcquiringStatus.ACQUIRING)
     falcon._pv_timeout = 0.1
-    with pytest.raises(TimeoutError):
-        falcon.on_stage()
+    with pytest.raises(StatusTimeoutError):
+        status = falcon.on_stage()
+        status.wait(timeout=2)
 
 
 def test_devices_falcon_unstage(falcon):
@@ -67,8 +69,9 @@ def test_devices_falcon_unstage(falcon):
     # Should timeout
     falcon.acquiring.put(FalconAcquiringStatus.ACQUIRING)
     falcon._pv_timeout = 0.1
-    with pytest.raises(TimeoutError):
-        falcon.on_unstage()
+    with pytest.raises(StatusTimeoutError):
+        status = falcon.on_unstage()
+        status.wait(timeout=2)
 
 
 def test_devices_falcon_stop(falcon):

tests/tests_devices/test_devices_trigger.py

@@ -9,7 +9,7 @@ from bec_server.device_server.tests.utils import DMMock
 from ophyd import DeviceStatus
 from ophyd_devices.tests.utils import MockPV, patch_dual_pvs
 
-from superxas_bec.devices.trigger import ContinuousSamplingMode, Trigger
+from superxas_bec.devices.trigger import ContinuousSamplingMode, SamplingDone, Trigger
 
 # pylint: disable=protected-access
 # pylint: disable=redefined-outer-name
@@ -94,13 +94,19 @@ def test_devices_trigger_trigger(trigger):
 
     trigger_status = DeviceStatus(device=trigger)
     trigger_status.set_finished()
-    with mock.patch.object(
-        trigger.task_handler, "submit_task", return_value=trigger_status
-    ) as mock_submit:
-        status = trigger.trigger()
-        assert falcon._stop_erase_and_wait_for_acquiring.call_count == 1
-        assert trigger.smpl.get() == 1  # smpl called with 1
-        # TODO check that the task_handler is called with the correct function
-        # This is currently not easily testable
-        assert mock_submit.call_count == 1
-        assert trigger_status == status
+    trigger.smpl_done._read_pv.mock_data = SamplingDone.DONE
+    status = trigger.trigger()
+    assert falcon._stop_erase_and_wait_for_acquiring.call_count == 1
+    assert trigger.smpl.get() == 1  # smpl called with 1
+    assert status.done is False
+    trigger.smpl_done._read_pv.mock_data = SamplingDone.RUNNING
+    assert status.done is False
+    trigger.smpl_done._read_pv.mock_data = SamplingDone.DONE
+    status.wait(timeout=2)
+    assert status.done is True
+    assert status.success is True
+
+    # # TODO check that the task_handler is called with the correct function
+    # # This is currently not easily testable
+    # assert mock_submit.call_count == 1
+    # assert trigger_status == status