Files
AareDAQ/tests/unit/daq/test_workflows.py
2026-07-06 15:39:02 +02:00

243 lines
7.3 KiB
Python

from unittest.mock import MagicMock
import pytest
from aare.daq.config import ABR_OMEGA_MOUNT, ABR_POS_MOUNT
from aare.daq.workflows import (
common2dh,
common_2rse,
dc2rse,
dc2sa,
dc2xtal_snapshot,
dh2sa,
sa2dc,
sa2dh,
sa2rse,
sa2se,
sa2xrf,
sa2xtal_snapshot,
se2sa,
xtal_snapshot2dc,
xtal_snapshot2sa,
)
from aare.devices.area_detector import AutoEnum
from aare.devices.bec_worker import BeamlineState
@pytest.fixture
def mock_devs():
devs = MagicMock()
devs.bec_worker = MagicMock()
devs.bec_worker.planner = MagicMock()
devs.bec_worker.move_to = MagicMock()
devs.dtz = None
return devs
@pytest.fixture
def mock_cfg():
cfg = MagicMock()
cfg.abr_meas_pos = "MEAS_POS"
cfg.dtz = None
cfg.dtz_safe_position = None
return cfg
def _assert_bec_moved(devs, state):
planner_calls = (
devs.bec_worker.planner.move_to.call_args_list if devs.bec_worker is not None else []
)
direct_calls = devs.bec_worker.move_to.call_args_list if devs.bec_worker is not None else []
assert any(call.args == (state,) for call in planner_calls) or any(
call.args == (state,) for call in direct_calls
), f"BEC was not asked to move to {state}"
def test_common_2rse(mock_devs, mock_cfg):
mock_devs.bec_worker = MagicMock()
mock_devs.bec_worker.planner = MagicMock()
mock_devs.bec_worker.move_to = MagicMock()
common_2rse(mock_devs, mock_cfg)
_assert_bec_moved(mock_devs, BeamlineState.ROBOT_SAMPLE_EXCHANGE)
mock_devs.smargon_move_home.assert_called_once()
assert mock_devs.aerotech_pos == ABR_POS_MOUNT
def test_common_2rse_moves_detector_to_safe_position_when_configured(mock_devs, mock_cfg):
mock_devs.dtz = 200
mock_cfg.dtz_safe_position = 600
common_2rse(mock_devs, mock_cfg)
mock_devs.bec_worker.det_z.assert_called_once_with(value=600)
_assert_bec_moved(mock_devs, BeamlineState.ROBOT_SAMPLE_EXCHANGE)
def test_common_2rse_ignores_missing_detector_safe_position(mock_devs, mock_cfg):
mock_devs.dtz = 200
mock_cfg.dtz_safe_position = None
common_2rse(mock_devs, mock_cfg)
mock_devs.bec_worker.det_z.assert_not_called()
_assert_bec_moved(mock_devs, BeamlineState.ROBOT_SAMPLE_EXCHANGE)
def test_dc2sa_moves_detector_to_safe_position_when_configured(mock_devs, mock_cfg):
mock_devs.dtz = 200
mock_cfg.dtz_safe_position = 600
dc2sa(mock_devs, mock_cfg)
mock_devs.bec_worker.det_z.assert_called_once_with(value=600)
_assert_bec_moved(mock_devs, BeamlineState.SAMPLE_ALIGNMENT)
assert mock_devs.aerotech_pos == mock_cfg.abr_meas_pos
def test_dc2sa_ignores_missing_detector_safe_position(mock_devs, mock_cfg):
mock_devs.dtz = 200
mock_cfg.dtz_safe_position = None
dc2sa(mock_devs, mock_cfg)
mock_devs.bec_worker.det_z.assert_not_called()
_assert_bec_moved(mock_devs, BeamlineState.SAMPLE_ALIGNMENT)
assert mock_devs.aerotech_pos == mock_cfg.abr_meas_pos
def test_sa2se(mock_devs, mock_cfg):
mock_devs.bec_worker = MagicMock()
mock_devs.bec_worker.planner = MagicMock()
mock_devs.bec_worker.move_to = MagicMock()
sa2se(mock_devs, mock_cfg)
mock_devs.samcam_auto.assert_any_call(AutoEnum.AUTO)
_assert_bec_moved(mock_devs, BeamlineState.MANUAL_SAMPLE_EXCHANGE)
assert mock_devs.aerotech_pos == ABR_POS_MOUNT
assert mock_devs.aerotech_omega == ABR_OMEGA_MOUNT
mock_devs.smargon_move_home.assert_called_once()
mock_devs.samcam_auto.assert_any_call(AutoEnum.ONCE)
def test_sa2rse(mock_devs, mock_cfg):
sa2rse(mock_devs, mock_cfg)
mock_devs.smargon_move_home.assert_called_once()
def test_sa2xtal_snapshot(mock_devs, mock_cfg):
mock_devs.bec_worker = MagicMock()
mock_devs.bec_worker.planner = MagicMock()
mock_devs.bec_worker.move_to = MagicMock()
sa2xtal_snapshot(mock_devs, mock_cfg)
assert mock_devs.samcam_settings.exposure == 0.001
_assert_bec_moved(mock_devs, BeamlineState.XTAL_SNAPSHOT)
def test_dc2xtal_snapshot(mock_devs, mock_cfg):
mock_devs.bec_worker = MagicMock()
mock_devs.bec_worker.planner = MagicMock()
mock_devs.bec_worker.move_to = MagicMock()
dc2xtal_snapshot(mock_devs, mock_cfg)
assert mock_devs.samcam_settings.exposure == 0.001
_assert_bec_moved(mock_devs, BeamlineState.XTAL_SNAPSHOT)
def test_xtal_snapshot2dc(mock_devs, mock_cfg):
mock_devs.bec_worker = MagicMock()
mock_devs.bec_worker.planner = MagicMock()
mock_devs.bec_worker.move_to = MagicMock()
xtal_snapshot2dc(mock_devs, mock_cfg)
_assert_bec_moved(mock_devs, BeamlineState.DATA_COLLECTION)
def test_xtal_snapshot2sa(mock_devs, mock_cfg):
mock_devs.bec_worker = MagicMock()
mock_devs.bec_worker.planner = MagicMock()
mock_devs.bec_worker.move_to = MagicMock()
xtal_snapshot2sa(mock_devs, mock_cfg)
_assert_bec_moved(mock_devs, BeamlineState.SAMPLE_ALIGNMENT)
def test_dc2rse(mock_devs, mock_cfg):
dc2rse(mock_devs, mock_cfg)
mock_devs.smargon_move_home.assert_called_once()
def test_se2sa(mock_devs, mock_cfg):
mock_devs.bec_worker = MagicMock()
mock_devs.bec_worker.planner = MagicMock()
mock_devs.bec_worker.move_to = MagicMock()
se2sa(mock_devs, mock_cfg)
_assert_bec_moved(mock_devs, BeamlineState.SAMPLE_ALIGNMENT)
assert mock_devs.aerotech_pos == mock_cfg.abr_meas_pos
def test_sa2dc(mock_devs, mock_cfg):
mock_devs.bec_worker = MagicMock()
mock_devs.bec_worker.planner = MagicMock()
mock_devs.bec_worker.move_to = MagicMock()
sa2dc(mock_devs, mock_cfg)
_assert_bec_moved(mock_devs, BeamlineState.DATA_COLLECTION)
def test_dc2sa(mock_devs, mock_cfg):
mock_devs.bec_worker = MagicMock()
mock_devs.bec_worker.planner = MagicMock()
mock_devs.bec_worker.move_to = MagicMock()
mock_cfg.dtz = None
dc2sa(mock_devs, mock_cfg)
_assert_bec_moved(mock_devs, BeamlineState.SAMPLE_ALIGNMENT)
assert mock_devs.aerotech_pos == mock_cfg.abr_meas_pos
def test_sa2xrf(mock_devs, mock_cfg):
mock_devs.bec_worker = MagicMock()
mock_devs.bec_worker.planner = MagicMock()
mock_devs.bec_worker.move_to = MagicMock()
sa2xrf(mock_devs, mock_cfg)
_assert_bec_moved(mock_devs, BeamlineState.FLUX_MEASUREMENT)
def test_sa2dh(mock_devs, mock_cfg):
mock_devs.tell.get_mounted_sample.return_value = "Sample"
sa2dh(mock_devs, mock_cfg)
mock_devs.tell.unmount.assert_called_with(wait=True)
mock_devs.tell.dry.assert_called_with(wait_cold=-1, wait=False)
def test_dh2sa(mock_devs, mock_cfg):
mock_devs.bec_worker = MagicMock()
mock_devs.bec_worker.planner = MagicMock()
mock_devs.bec_worker.move_to = MagicMock()
dh2sa(mock_devs, mock_cfg)
_assert_bec_moved(mock_devs, BeamlineState.SAMPLE_ALIGNMENT)
def test_common2dh_skips_dry_when_tell_already_in_ppark(mock_devs, mock_cfg):
mock_devs.bec_worker = MagicMock()
mock_devs.bec_worker.planner = MagicMock()
mock_devs.bec_worker.move_to = MagicMock()
mock_devs.tell = MagicMock()
mock_devs.tell.is_position.return_value = True
mock_devs.tell.get_mounted_sample.return_value = None
common2dh(mock_devs, mock_cfg)
mock_devs.tell.is_position.assert_called_once_with("pPark")
mock_devs.tell.get_mounted_sample.assert_called_once_with()
mock_devs.tell.unmount.assert_not_called()
mock_devs.tell.dry.assert_not_called()