diff --git a/slac_devices/magnet.py b/slac_devices/magnet.py index 4cc0363..41b5f58 100644 --- a/slac_devices/magnet.py +++ b/slac_devices/magnet.py @@ -1,5 +1,6 @@ from datetime import datetime from functools import wraps +import time from pydantic import ( Field, PositiveFloat, @@ -77,6 +78,12 @@ def __init__(self, *args, **kwargs): super(MagnetMetadata, self).__init__(*args, **kwargs) +class MagnetTimeoutError(Exception): + """Raised when the magnet fails to settle within the allowed time""" + + pass + + class Magnet(Device): controls_information: SerializeAsAny[MagnetControlInformation] metadata: SerializeAsAny[MagnetMetadata] @@ -169,7 +176,7 @@ def bdes(self) -> float: @bdes.setter def bdes(self, bval) -> None: - self.controls_information.PVs.bdes.put(value=bval) + self.controls_information.PVs.bdes.put(value=bval, wait=True) # wait until PV has been set @property def ctrl(self) -> str: @@ -203,8 +210,11 @@ def safe_level(self) -> float: def read_tolerance(self) -> float: return self.metadata.read_tolerance - def is_bact_settled(self, b_tolerance: Optional[float] = 0.0) -> bool: - return abs(self.bdes) - abs(self.bact) < b_tolerance + def is_bact_settled(self, b_tolerance: Optional[float] = 0.001) -> bool: + if b_tolerance: + return abs(self.bdes - self.bact) < b_tolerance + else: + return abs(self.bdes - self.bact) < 0.001 @check_options("TRIM") @check_state @@ -282,9 +292,24 @@ def scan( function: Optional[callable] = None, ) -> None: for setting in scan_settings: - self.bctrl = setting + self.set_bdes_with_validation(setting) function() if function else None + def set_bdes_with_validation(self, bval: float, settle_timeout_in_seconds: int = 5): + self.bdes = bval + self.trim() + time_when_trim_started = datetime.now() + while not self.is_bact_settled(self.b_tolerance): + # Timeout if magnet takes too long to settle + if ( + datetime.now() - time_when_trim_started + ).seconds > settle_timeout_in_seconds: + raise MagnetTimeoutError( + f"Took more than {settle_timeout_in_seconds} seconds for " + f"{self.name}:BACT to reach {self.name}:BDES." + ) + time.sleep(0.1) + class MagnetCollection(DeviceCollection): devices: Dict[str, SerializeAsAny[Magnet]] = Field(alias="magnets") diff --git a/tests/test_magnet.py b/tests/test_magnet.py index e4f8d7d..9b0d8f0 100644 --- a/tests/test_magnet.py +++ b/tests/test_magnet.py @@ -170,7 +170,7 @@ def test_bdes(self, mock_pv_get) -> None: def test_set_bdes(self, mock_pv_put) -> None: mock_pv_put.return_value = None self.magnet.bdes = 0.1 - mock_pv_put.assert_called_once_with(value=0.1) + mock_pv_put.assert_called_once_with(value=0.1, wait=True) @patch("epics.PV.get", new_callable=Mock) def test_get_bctrl(self, mock_pv_get) -> None: @@ -399,7 +399,7 @@ def test_set_bdes_with_args( } mock_bact_settled.return_value = True self.magnet_collection.set_bdes(magnet_dict=bdes_settings) - mock_bdes_put.assert_called_once_with(value=0.1) + mock_bdes_put.assert_called_once_with(value=0.1, wait=True) mock_trim.assert_called_once() mock_bact_settled.assert_called_once_with()