Move WarningException and some oav plans to dodal (#1423)

* Move WarningException and some oav plans to dodal * make wrapper function to raise WarningExceptions from plans * Simplify xrc tests slightly --------- Co-authored-by: Dominic Oram <[email protected]>
DiamondLightSource · Jun 12, 2024 · c5722de · c5722de
1 parent 35a6910
commit c5722de
Show file tree

Hide file tree

Showing 10 changed files with 88 additions and 181 deletions.
diff --git a/setup.cfg b/setup.cfg
@@ -45,7 +45,7 @@ install_requires =
     ophyd-async >= 0.3a5
     bluesky >= 1.13.0a3
     blueapi >= 0.4.3-rc1
-    dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@8e62a4db4fc07e1d873f84a155c85105a986649b
+    dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@ce2a48780b979fb585214092741b76c449b05452
 
 [options.entry_points]
 console_scripts =

diff --git a/src/hyperion/device_setup_plans/setup_oav.py b/src/hyperion/device_setup_plans/setup_oav.py
@@ -1,20 +1,12 @@
 from functools import partial
-from typing import Generator, Tuple
 
 import bluesky.plan_stubs as bps
-import numpy as np
-from bluesky.utils import Msg
-from dodal.devices.oav.oav_calculations import camera_coordinates_to_xyz
-from dodal.devices.oav.oav_detector import OAV, OAVConfigParams
+from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.oav.oav_errors import OAVError_ZoomLevelNotFound
 from dodal.devices.oav.oav_parameters import OAVParameters
 from dodal.devices.oav.pin_image_recognition import PinTipDetection
 from dodal.devices.oav.utils import ColorMode
-from dodal.devices.smargon import Smargon
 
-from hyperion.exceptions import WarningException
-
-Pixel = Tuple[int, int]
 oav_group = "oav_setup"
 # Helper function to make sure we set the waiting groups correctly
 set_using_group = partial(bps.abs_set, group=oav_group)
@@ -94,49 +86,3 @@ def pre_centring_setup_oav(
     """
     TODO: We require setting the backlight brightness to that in the json, we can't do this currently without a PV.
     """
-
-
-def calculate_x_y_z_of_pixel(
-    current_x_y_z, current_omega, pixel: Pixel, oav_params: OAVConfigParams
-) -> np.ndarray:
-    beam_distance_px: Pixel = oav_params.calculate_beam_distance(*pixel)
-
-    assert oav_params.micronsPerXPixel
-    assert oav_params.micronsPerYPixel
-    return current_x_y_z + camera_coordinates_to_xyz(
-        beam_distance_px[0],
-        beam_distance_px[1],
-        current_omega,
-        oav_params.micronsPerXPixel,
-        oav_params.micronsPerYPixel,
-    )
-
-
-def get_move_required_so_that_beam_is_at_pixel(
-    smargon: Smargon, pixel: Pixel, oav_params: OAVConfigParams
-) -> Generator[Msg, None, np.ndarray]:
-    """Calculate the required move so that the given pixel is in the centre of the beam."""
-
-    current_motor_xyz = np.array(
-        [
-            (yield from bps.rd(smargon.x)),
-            (yield from bps.rd(smargon.y)),
-            (yield from bps.rd(smargon.z)),
-        ],
-        dtype=np.float64,
-    )
-    current_angle = yield from bps.rd(smargon.omega)
-
-    return calculate_x_y_z_of_pixel(current_motor_xyz, current_angle, pixel, oav_params)
-
-
-def wait_for_tip_to_be_found(
-    ophyd_pin_tip_detection: PinTipDetection,
-) -> Generator[Msg, None, Pixel]:
-    yield from bps.trigger(ophyd_pin_tip_detection, wait=True)
-    found_tip = yield from bps.rd(ophyd_pin_tip_detection.triggered_tip)
-    if found_tip == ophyd_pin_tip_detection.INVALID_POSITION:
-        timeout = yield from bps.rd(ophyd_pin_tip_detection.validity_timeout)
-        raise WarningException(f"No pin found after {timeout} seconds")
-
-    return found_tip  # type: ignore
diff --git a/src/hyperion/exceptions.py b/src/hyperion/exceptions.py
@@ -1,5 +1,46 @@
+from typing import Callable, Generator, Type, TypeVar
+
+from bluesky.plan_stubs import null
+from bluesky.preprocessors import contingency_wrapper
+from bluesky.utils import Msg
+
+
 class WarningException(Exception):
     """An exception used when we want to warn GDA of a
     problem but continue with UDC anyway"""
 
     pass
+
+
+T = TypeVar("T")
+
+
+def catch_exception_and_warn(
+    exception_to_catch: Type[Exception],
+    func: Callable[..., Generator[Msg, None, T]],
+    *args,
+    **kwargs,
+) -> Generator[Msg, None, T]:
+    """A plan wrapper to catch a specific exception and instead raise a WarningException,
+    so that UDC is not halted
+
+    Example usage:
+
+    'def plan_which_can_raise_exception_a(*args, **kwargs):
+        ...
+    yield from catch_exception_and_warn(ExceptionA, plan_which_can_raise_exception_a, **args, **kwargs)'
+
+    This will catch ExceptionA raised by the plan and instead raise a WarningException
+    """
+
+    def warn_if_exception_matches(exception: Exception):
+        if isinstance(exception, exception_to_catch):
+            raise WarningException(str(exception))
+        yield from null()
+
+    return (
+        yield from contingency_wrapper(
+            func(*args, **kwargs),
+            except_plan=warn_if_exception_matches,
+        )
+    )
diff --git a/src/hyperion/experiment_plans/oav_grid_detection_plan.py b/src/hyperion/experiment_plans/oav_grid_detection_plan.py
@@ -11,12 +11,13 @@
 from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.oav.pin_image_recognition import PinTipDetection
 from dodal.devices.oav.pin_image_recognition.utils import NONE_VALUE
+from dodal.devices.oav.utils import PinNotFoundException, wait_for_tip_to_be_found
 from dodal.devices.smargon import Smargon
 
 from hyperion.device_setup_plans.setup_oav import (
     pre_centring_setup_oav,
-    wait_for_tip_to_be_found,
 )
+from hyperion.exceptions import catch_exception_and_warn
 from hyperion.log import LOGGER
 from hyperion.parameters.constants import CONST
 from hyperion.utils.context import device_composite_from_context
@@ -103,7 +104,9 @@ def grid_detection_plan(
         # See #673 for improvements
         yield from bps.sleep(CONST.HARDWARE.OAV_REFRESH_DELAY)
 
-        tip_x_px, tip_y_px = yield from wait_for_tip_to_be_found(pin_tip_detection)
+        tip_x_px, tip_y_px = yield from catch_exception_and_warn(
+            PinNotFoundException, wait_for_tip_to_be_found, pin_tip_detection
+        )
 
         LOGGER.info(f"Tip is at x,y: {tip_x_px},{tip_y_px}")
 

diff --git a/src/hyperion/experiment_plans/pin_tip_centring_plan.py b/src/hyperion/experiment_plans/pin_tip_centring_plan.py
@@ -9,14 +9,14 @@
 from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.oav.oav_parameters import OAV_CONFIG_JSON, OAVParameters
 from dodal.devices.oav.pin_image_recognition import PinTipDetection
-from dodal.devices.smargon import Smargon
-
-from hyperion.device_setup_plans.setup_oav import (
+from dodal.devices.oav.utils import (
     Pixel,
     get_move_required_so_that_beam_is_at_pixel,
-    pre_centring_setup_oav,
     wait_for_tip_to_be_found,
 )
+from dodal.devices.smargon import Smargon
+
+from hyperion.device_setup_plans.setup_oav import pre_centring_setup_oav
 from hyperion.exceptions import WarningException
 from hyperion.log import LOGGER
 from hyperion.parameters.constants import CONST

diff --git a/src/hyperion/external_interaction/callbacks/grid_detection_callback.py b/src/hyperion/external_interaction/callbacks/grid_detection_callback.py
@@ -3,9 +3,9 @@
 import numpy as np
 from bluesky.callbacks import CallbackBase
 from dodal.devices.oav.oav_detector import OAVConfigParams
+from dodal.devices.oav.utils import calculate_x_y_z_of_pixel
 from event_model.documents import Event
 
-from hyperion.device_setup_plans.setup_oav import calculate_x_y_z_of_pixel
 from hyperion.log import LOGGER
 
 

diff --git a/tests/unit_tests/device_setup_plans/test_setup_oav.py b/tests/unit_tests/device_setup_plans/test_setup_oav.py
@@ -1,26 +1,18 @@
-from functools import partial
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import MagicMock
 
-import numpy as np
 import pytest
 from bluesky import plan_stubs as bps
 from bluesky.run_engine import RunEngine
 from dodal.beamlines import i03
 from dodal.devices.oav.oav_detector import OAV, OAVConfigParams
 from dodal.devices.oav.oav_parameters import OAVParameters
 from dodal.devices.oav.pin_image_recognition import PinTipDetection
-from dodal.devices.oav.pin_image_recognition.utils import SampleLocation
-from dodal.devices.smargon import Smargon
 from ophyd.signal import Signal
-from ophyd.sim import instantiate_fake_device
 from ophyd.status import Status
 
 from hyperion.device_setup_plans.setup_oav import (
-    get_move_required_so_that_beam_is_at_pixel,
     pre_centring_setup_oav,
-    wait_for_tip_to_be_found,
 )
-from hyperion.exceptions import WarningException
 
 ZOOM_LEVELS_XML = "tests/test_data/test_jCameraManZoomLevels.xml"
 OAV_CENTRING_JSON = "tests/test_data/test_OAVCentring.json"
@@ -49,31 +41,6 @@ def mock_parameters():
     return OAVParameters("loopCentring", OAV_CENTRING_JSON)
 
 
-def fake_smargon() -> Smargon:
-    smargon = i03.smargon(fake_with_ophyd_sim=True)
-    smargon.x.user_setpoint._use_limits = False
-    smargon.y.user_setpoint._use_limits = False
-    smargon.z.user_setpoint._use_limits = False
-    smargon.omega.user_setpoint._use_limits = False
-
-    def mock_set(motor, val):
-        motor.user_readback.sim_put(val)  # type: ignore
-        return Status(done=True, success=True)
-
-    def patch_motor(motor):
-        return patch.object(motor, "set", partial(mock_set, motor))
-
-    with patch_motor(smargon.omega), patch_motor(smargon.x), patch_motor(
-        smargon.y
-    ), patch_motor(smargon.z):
-        return smargon
-
-
-@pytest.fixture
-def smargon():
-    yield fake_smargon()
-
-
 @pytest.mark.parametrize(
     "zoom, expected_plugin",
     [
@@ -95,47 +62,6 @@ def test_when_set_up_oav_with_different_zoom_levels_then_flat_field_applied_corr
     assert oav.grid_snapshot.input_plugin.get() == expected_plugin
 
 
-@pytest.mark.parametrize(
-    "px_per_um, beam_centre, angle, pixel_to_move_to, expected_xyz",
-    [
-        # Simple case of beam being in the top left and each pixel being 1 mm
-        ([1000, 1000], [0, 0], 0, [100, 190], [100, 190, 0]),
-        ([1000, 1000], [0, 0], -90, [50, 250], [50, 0, 250]),
-        ([1000, 1000], [0, 0], 90, [-60, 450], [-60, 0, -450]),
-        # Beam offset
-        ([1000, 1000], [100, 100], 0, [100, 100], [0, 0, 0]),
-        ([1000, 1000], [100, 100], -90, [50, 250], [-50, 0, 150]),
-        # Pixels_per_micron different
-        ([10, 50], [0, 0], 0, [100, 190], [1, 9.5, 0]),
-        ([60, 80], [0, 0], -90, [50, 250], [3, 0, 20]),
-    ],
-)
-def test_values_for_move_so_that_beam_is_at_pixel(
-    smargon: Smargon,
-    oav: OAV,
-    px_per_um,
-    beam_centre,
-    angle,
-    pixel_to_move_to,
-    expected_xyz,
-):
-    oav.parameters.micronsPerXPixel = px_per_um[0]
-    oav.parameters.micronsPerYPixel = px_per_um[1]
-    oav.parameters.beam_centre_i = beam_centre[0]
-    oav.parameters.beam_centre_j = beam_centre[1]
-
-    smargon.omega.user_readback.sim_put(angle)  # type: ignore
-
-    RE = RunEngine(call_returns_result=True)
-    pos = RE(
-        get_move_required_so_that_beam_is_at_pixel(
-            smargon, pixel_to_move_to, oav.parameters
-        )
-    ).plan_result
-
-    assert pos == pytest.approx(expected_xyz)
-
-
 def test_when_set_up_oav_then_only_waits_on_oav_to_finish(
     mock_parameters: OAVParameters, oav: OAV, ophyd_pin_tip_detection: PinTipDetection
 ):
@@ -150,35 +76,3 @@ def my_plan():
 
     RE = RunEngine()
     RE(my_plan())
-
-
-@pytest.mark.asyncio
-async def test_given_tip_found_when_wait_for_tip_to_be_found_called_then_tip_immediately_returned():
-    mock_pin_tip_detect: PinTipDetection = instantiate_fake_device(
-        PinTipDetection, name="pin_detect"
-    )
-    await mock_pin_tip_detect.connect(mock=True)
-    mock_pin_tip_detect._get_tip_and_edge_data = AsyncMock(
-        return_value=SampleLocation(100, 100, np.array([]), np.array([]))
-    )
-    RE = RunEngine(call_returns_result=True)
-    result = RE(wait_for_tip_to_be_found(mock_pin_tip_detect))
-    assert result.plan_result == (100, 100)  # type: ignore
-    mock_pin_tip_detect._get_tip_and_edge_data.assert_called_once()
-
-
-@pytest.mark.asyncio
-async def test_given_no_tip_when_wait_for_tip_to_be_found_called_then_exception_thrown():
-    mock_pin_tip_detect: PinTipDetection = instantiate_fake_device(
-        PinTipDetection, name="pin_detect"
-    )
-    await mock_pin_tip_detect.connect(mock=True)
-    await mock_pin_tip_detect.validity_timeout.set(0.2)
-    mock_pin_tip_detect._get_tip_and_edge_data = AsyncMock(
-        return_value=SampleLocation(
-            *PinTipDetection.INVALID_POSITION, np.array([]), np.array([])
-        )
-    )
-    RE = RunEngine(call_returns_result=True)
-    with pytest.raises(WarningException):
-        RE(wait_for_tip_to_be_found(mock_pin_tip_detect))