From 8ea72b070ef43b9225e95fd4c6f0448c0315c3d0 Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Thu, 19 Sep 2024 18:49:29 +0100 Subject: [PATCH] Revert zoom level after thawing --- src/mx_bluesky/beamlines/i04/thawing_plan.py | 14 ++- tests/conftest.py | 9 ++ .../unit_tests/beamlines/i04/test_thawing.py | 92 ++++++++++++++++--- 3 files changed, 100 insertions(+), 15 deletions(-) diff --git a/src/mx_bluesky/beamlines/i04/thawing_plan.py b/src/mx_bluesky/beamlines/i04/thawing_plan.py index aab664f34..549d05e84 100644 --- a/src/mx_bluesky/beamlines/i04/thawing_plan.py +++ b/src/mx_bluesky/beamlines/i04/thawing_plan.py @@ -24,10 +24,9 @@ def thaw_and_stream_to_redis( ) -> MsgGenerator: zoom_percentage = yield from bps.rd(oav.zoom_controller.percentage) sample_id = yield from bps.rd(robot.sample_id) + zoom_level_before_thawing = yield from bps.rd(oav.zoom_controller.level) yield from bps.mv(oav.zoom_controller.level, "1.0x") - yield from bps.mv(oav_to_redis_forwarder.sample_id, sample_id) - yield from bps.mv(oav_to_redis_forwarder.selected_source, Source.FULL_SCREEN) def switch_forwarder_to_ROI(): yield from bps.complete(oav_to_redis_forwarder) @@ -46,6 +45,9 @@ def switch_forwarder_to_ROI(): } ) def _thaw_and_stream_to_redis(): + yield from bps.mv(oav_to_redis_forwarder.sample_id, sample_id) + yield from bps.mv(oav_to_redis_forwarder.selected_source, Source.FULL_SCREEN) + yield from bps.kickoff(oav_to_redis_forwarder, wait=True) yield from bps.monitor(smargon.omega.user_readback, name="smargon") yield from bps.monitor(oav_to_redis_forwarder.uuid, name="oav") @@ -54,7 +56,13 @@ def _thaw_and_stream_to_redis(): ) yield from bps.complete(oav_to_redis_forwarder) - yield from _thaw_and_stream_to_redis() + def cleanup(): + yield from bps.mv(oav.zoom_controller.level, zoom_level_before_thawing) + + yield from bpp.contingency_wrapper( + _thaw_and_stream_to_redis(), + final_plan=cleanup, + ) def thaw( diff --git a/tests/conftest.py b/tests/conftest.py index 3bb0ddaf4..c4380a9f6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -346,6 +346,15 @@ def oav(test_config_files): test_config_files["zoom_params_file"], test_config_files["display_config"] ) oav = i03.oav(fake_with_ophyd_sim=True, params=parameters) + + oav.zoom_controller.zrst.set("1.0x") + oav.zoom_controller.onst.set("2.0x") + + oav.parameters.micronsPerXPixel = 1.58 + oav.parameters.micronsPerYPixel = 1.58 + oav.parameters.beam_centre_i = 517 + oav.parameters.beam_centre_j = 350 + oav.snapshot.trigger = MagicMock(return_value=NullStatus()) oav.grid_snapshot.trigger = MagicMock(return_value=NullStatus()) return oav diff --git a/tests/unit_tests/beamlines/i04/test_thawing.py b/tests/unit_tests/beamlines/i04/test_thawing.py index e9b26a5ff..b909ec515 100644 --- a/tests/unit_tests/beamlines/i04/test_thawing.py +++ b/tests/unit_tests/beamlines/i04/test_thawing.py @@ -1,4 +1,5 @@ from collections.abc import AsyncGenerator +from functools import partial from unittest.mock import ANY, MagicMock, call, patch import pytest @@ -11,7 +12,7 @@ from dodal.devices.robot import BartRobot from dodal.devices.smargon import Smargon from dodal.devices.thawer import Thawer, ThawerStates -from ophyd.sim import NullStatus, instantiate_fake_device +from ophyd.sim import NullStatus from ophyd_async.core import ( AsyncStatus, DeviceCollector, @@ -31,17 +32,6 @@ class MyException(Exception): pass -@pytest.fixture -def oav() -> OAV: - oav: OAV = instantiate_fake_device(OAV, params=MagicMock()) - - oav.zoom_controller.zrst.set("1.0x") - - oav.wait_for_connection() - - return oav - - def patch_motor(motor: Motor, initial_position: float = 0): set_mock_value(motor.user_setpoint, initial_position) set_mock_value(motor.user_readback, initial_position) @@ -290,3 +280,81 @@ def test_thaw_and_stream_will_switch_murko_source_half_way_through_thaw( lambda msg: msg.command == "complete" and msg.obj.name == "oav_to_redis_forwarder", ) + + +def _run_thaw_and_stream_and_assert_zoom_changes( + smargon: Smargon, + thawer: Thawer, + oav_forwarder: OAVToRedisForwarder, + oav: OAV, + RE: RunEngine, + expect_raises=None, +): + mock_set = MagicMock() + + def cb(*args, value, **kwargs): + mock_set(value) + + oav.zoom_controller.level.sim_put("2.0x") # type:ignore + oav.zoom_controller.level.subscribe(cb) + + run_plan = partial( + RE, + thaw_and_stream_to_redis( + 10, + 360, + thawer=thawer, + smargon=smargon, + oav=oav, + robot=MagicMock(), + oav_to_redis_forwarder=oav_forwarder, + ), + ) + + if expect_raises: + with pytest.raises(expect_raises): + run_plan() + else: + run_plan() + + mock_set.assert_has_calls( + [ + call( + "1.0x", + ), + call( + "2.0x", + ), + ] + ) + + +@patch("mx_bluesky.beamlines.i04.thawing_plan.MurkoCallback") +def test_given_thaw_succeeds_then_thaw_and_stream_sets_zoom_to_1_and_back( + patch_murko_callback, + smargon: Smargon, + thawer: Thawer, + oav_forwarder: OAVToRedisForwarder, + oav: OAV, + RE: RunEngine, +): + _run_thaw_and_stream_and_assert_zoom_changes( + smargon, thawer, oav_forwarder, oav, RE + ) + + +@patch("mx_bluesky.beamlines.i04.thawing_plan.MurkoCallback") +@patch("mx_bluesky.beamlines.i04.thawing_plan.thaw") +def test_given_thaw_fails_then_thaw_and_stream_sets_zoom_to_1_and_back( + mock_thaw, + patch_murko_callback, + smargon: Smargon, + thawer: Thawer, + oav_forwarder: OAVToRedisForwarder, + oav: OAV, + RE: RunEngine, +): + mock_thaw.side_effect = Exception() + _run_thaw_and_stream_and_assert_zoom_changes( + smargon, thawer, oav_forwarder, oav, RE, Exception + )