Revert zoom level after thawing

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(

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

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
+ )