AareDAQ/tests/unit/daq/test_workflows.py

import pytest
from unittest.mock import MagicMock, patch
from aare.daq.workflows import move_bsz, common_2rse, sa2se, sa2rse, sa2xtal_snapshot, dc2xtal_snapshot, xtal_snapshot2dc, xtal_snapshot2sa, dc2rse, se2sa, sa2dc, dc2sa, sa2xrf, sa2dh, dh2sa
from aare.daq.config import ABR_POS_MOUNT
from aare.common.models import StagePositionEnum
from aare.devices.area_detector import AutoEnum
from aare.devices.bec_worker import BeamlineState


@pytest.fixture
def mock_devs():
    devs = MagicMock()
    devs.bsz.position = 0.0
    devs.beamstop_stage_up = False
    devs.reflector_up = False
    devs.bec_worker = MagicMock()
    devs.bec_worker.planner = MagicMock()
    devs.bec_worker.move_to = MagicMock()
    return devs


@pytest.fixture
def mock_cfg():
    cfg = MagicMock()
    cfg.abr_meas_pos = "MEAS_POS"
    cfg.dtz = 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_move_bsz_no_move(mock_devs):
    mock_devs.bsz.position = 1.0
    move_bsz(mock_devs, 1.05)
    assert mock_devs.beamstop_z.call_count == 0


def test_move_bsz_with_move(mock_devs):
    mock_devs.bsz.position = 0.0
    mock_devs.beamstop_stage_up = False
    mock_devs.reflector_up = True

    move_bsz(mock_devs, 1.0)

    assert mock_devs.reflector_up is True
    assert mock_devs.beamstop_stage_up is False
    assert mock_devs.beamstop_z == 1.0


def test_common_2rse_no_bec(mock_devs, mock_cfg):
    mock_devs.bec_worker = None
    common_2rse(mock_devs, mock_cfg)

    assert mock_devs.collimator == 20.0
    assert mock_devs.scintillator == 20.0
    mock_devs.smargon_move_home.assert_called_once()
    assert mock_devs.beamstop_stage_up == StagePositionEnum.PARK


def test_common_2rse_with_bec(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_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 == mock_cfg.abr_meas_pos
    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.0005
    _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.0005
    _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)