AareDAQ/tests/unit/daq/test_workflows.py

from types import SimpleNamespace

import pytest
from unittest.mock import MagicMock, patch
from aare.daq.workflows import common_2rse, sa2se, sa2rse, sa2xtal_snapshot, dc2xtal_snapshot, xtal_snapshot2dc, xtal_snapshot2sa, dc2rse, se2sa, sa2dc, dc2sa, sa2xrf, sa2dh, dh2sa, common2dh
from aare.daq.config import ABR_POS_MOUNT, ABR_OMEGA_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.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
    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_cfg.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_cfg.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_cfg.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_cfg.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()