tests/tests_devices/test_devices_falcon.py

@@ -5,8 +5,8 @@ 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 ophyd_devices.utils.psi_device_base_utils import StatusTimeoutErrorWithErrorInfo
 
 from superxas_bec.devices.falcon import FalconAcquiringStatus, FalconSuperXAS
 
@@ -51,7 +51,7 @@ def test_devices_falcon_stage(falcon):
     # Should timeout
     falcon.acquiring.put(FalconAcquiringStatus.ACQUIRING)
     falcon._pv_timeout = 0.1
-    with pytest.raises(StatusTimeoutError):
+    with pytest.raises(StatusTimeoutErrorWithErrorInfo):
         status = falcon.on_stage()
         status.wait(timeout=2)
 
@@ -69,7 +69,7 @@ def test_devices_falcon_unstage(falcon):
     # Should timeout
     falcon.acquiring.put(FalconAcquiringStatus.ACQUIRING)
     falcon._pv_timeout = 0.1
-    with pytest.raises(StatusTimeoutError):
+    with pytest.raises(StatusTimeoutErrorWithErrorInfo):
         status = falcon.on_unstage()
         status.wait(timeout=2)