bec/scan_server/tests/test_scan_guard.py

from unittest import mock

import pytest
from bec_lib import MessageEndpoints, messages
from bec_lib.redis_connector import MessageObject
from bec_lib.tests.utils import dm, dm_with_devices
from utils import scan_server_mock

from scan_server.scan_guard import ScanGuard, ScanRejection, ScanStatus


@pytest.fixture
def scan_guard_mock(scan_server_mock):
    sg = ScanGuard(parent=scan_server_mock)
    sg.device_manager.connector = mock.MagicMock()
    yield sg


@pytest.mark.parametrize(
    "scan_queue_msg",
    [
        (
            messages.ScanQueueMessage(
                scan_type="fermat_scan",
                parameter={"args": {"samx": (-5, 5), "samy": (-5, 5)}, "kwargs": {"step": 3}},
                queue="primary",
            )
        ),
        (
            messages.ScanQueueMessage(
                scan_type="device_rpc",
                parameter={"device": "samy", "args": {}, "kwargs": {}},
                queue="primary",
            )
        ),
        (
            messages.ScanQueueMessage(
                scan_type="device_rpc",
                parameter={"device": ["samy"], "args": {}, "kwargs": {}},
                queue="primary",
            )
        ),
    ],
)
def test_check_motors_movable_enabled(scan_server_mock, scan_queue_msg):
    k = scan_server_mock

    sg = ScanGuard(parent=k)
    sg._check_motors_movable(scan_queue_msg)
    config_reply = messages.RequestResponseMessage(accepted=True, message="")
    with mock.patch.object(
        k.device_manager.config_helper, "wait_for_config_reply", return_value=config_reply
    ):
        with mock.patch.object(k.device_manager.config_helper, "wait_for_service_response"):
            k.device_manager.devices["samx"].enabled = True
            k.device_manager.devices["samy"].enabled = False
            with pytest.raises(ScanRejection) as scan_rejection:
                sg._check_motors_movable(scan_queue_msg)
            assert "Device samy is not enabled." in scan_rejection.value.args


@pytest.mark.parametrize("device,func,is_valid", [("samx", "read", True)])
def test_device_rpc_is_valid(scan_guard_mock, device, func, is_valid):
    sg = scan_guard_mock
    assert sg._device_rpc_is_valid(device, func) == is_valid


@pytest.mark.parametrize(
    "scan_queue_msg,valid",
    [
        (
            messages.ScanQueueMessage(
                scan_type="fermat_scan",
                parameter={"args": {"samx": (-5, 5), "samy": (-5, 5)}, "kwargs": {"step": 3}},
                queue="primary",
            ),
            True,
        ),
        (
            messages.ScanQueueMessage(
                scan_type="device_rpc",
                parameter={"device": "samy", "args": {}, "kwargs": {}},
                queue="primary",
            ),
            True,
        ),
        (
            messages.ScanQueueMessage(
                scan_type="device_rpc",
                parameter={"device": ["samy"], "args": {}, "kwargs": {}},
                queue="primary",
            ),
            True,
        ),
    ],
)
def test_valid_request(scan_server_mock, scan_queue_msg, valid):
    k = scan_server_mock

    sg = ScanGuard(parent=k)
    config_reply = messages.RequestResponseMessage(accepted=True, message="")
    with mock.patch.object(
        k.device_manager.config_helper, "wait_for_config_reply", return_value=config_reply
    ):
        with mock.patch.object(k.device_manager.config_helper, "wait_for_service_response"):
            with mock.patch.object(sg, "_check_valid_scan") as valid_scan:
                k.device_manager.devices["samx"].enabled = True
                k.device_manager.devices["samy"].enabled = True
                status = sg._is_valid_scan_request(scan_queue_msg)
                valid_scan.assert_called_once_with(scan_queue_msg)
                assert status.accepted == valid


def test_check_valid_scan_raises_for_unknown_scan(scan_guard_mock):
    sg = scan_guard_mock
    sg.connector = mock.MagicMock()
    sg.connector.get.return_value = messages.AvailableResourceMessage(
        resource={"fermat_scan": "fermat_scan"}
    )

    request = messages.ScanQueueMessage(
        scan_type="unknown_scan",
        parameter={"args": {"samx": (-5, 5), "samy": (-5, 5)}, "kwargs": {"step": 3}},
        queue="primary",
    )

    with pytest.raises(ScanRejection) as scan_rejection:
        sg._check_valid_scan(request)


def test_check_valid_scan_accepts_known_scan(scan_guard_mock):
    sg = scan_guard_mock
    sg.connector = mock.MagicMock()
    sg.connector.get.return_value = messages.AvailableResourceMessage(
        resource={"fermat_scan": "fermat_scan"}
    )

    request = messages.ScanQueueMessage(
        scan_type="fermat_scan",
        parameter={"args": {"samx": (-5, 5), "samy": (-5, 5)}, "kwargs": {"step": 3}},
        queue="primary",
    )

    sg._check_valid_scan(request)


def test_check_valid_scan_device_rpc(scan_guard_mock):
    sg = scan_guard_mock
    sg.connector = mock.MagicMock()
    sg.connector.get.return_value = messages.AvailableResourceMessage(
        resource={"device_rpc": "device_rpc"}
    )
    request = messages.ScanQueueMessage(
        scan_type="device_rpc",
        parameter={"device": "samy", "func": "read", "args": {}, "kwargs": {}},
        queue="primary",
    )
    with mock.patch.object(sg, "_device_rpc_is_valid") as rpc_valid:
        sg._check_valid_scan(request)
        rpc_valid.assert_called_once_with(device="samy", func="read")


def test_check_valid_scan_device_rpc_raises(scan_guard_mock):
    sg = scan_guard_mock
    sg.connector = mock.MagicMock()
    sg.connector.get.return_value = messages.AvailableResourceMessage(
        resource={"device_rpc": "device_rpc"}
    )
    request = messages.ScanQueueMessage(
        scan_type="device_rpc",
        parameter={"device": "samy", "func": "read", "args": {}, "kwargs": {}},
        queue="primary",
    )
    with pytest.raises(ScanRejection) as scan_rejection:
        with mock.patch.object(sg, "_device_rpc_is_valid") as rpc_valid:
            rpc_valid.return_value = False
            sg._check_valid_scan(request)
            rpc_valid.assert_called_once_with(device="samy", func="read")
        assert "Rejected rpc: " in scan_rejection.value.args


def test_handle_scan_modification_request(scan_guard_mock):
    sg = scan_guard_mock
    msg = messages.ScanQueueModificationMessage(
        scanID="scanID", action="abort", parameter={}, metadata={"RID": "RID"}
    )
    with mock.patch.object(sg.device_manager.connector, "send") as send:
        sg._handle_scan_modification_request(msg)
        send.assert_called_once_with(MessageEndpoints.scan_queue_modification(), msg)


def test_handle_scan_modification_request_restart(scan_guard_mock):
    sg = scan_guard_mock
    msg = messages.ScanQueueModificationMessage(
        scanID="scanID", action="restart", parameter={"RID": "RID"}, metadata={"RID": "new_RID"}
    )
    with mock.patch.object(sg, "_send_scan_request_response") as send_response:
        with mock.patch("scan_server.scan_guard.ScanStatus") as scan_status:
            sg._handle_scan_modification_request(msg)
            send_response.assert_called_once_with(scan_status(), {"RID": "RID"})


def test_append_to_scan_queue(scan_guard_mock):
    sg = scan_guard_mock
    msg = messages.ScanQueueMessage(
        scan_type="fermat_scan",
        parameter={"args": {"samx": (-5, 5), "samy": (-5, 5)}, "kwargs": {"step": 3}},
        queue="primary",
    )
    with mock.patch.object(sg.device_manager.connector, "send") as send:
        sg._append_to_scan_queue(msg)
        send.assert_called_once_with(MessageEndpoints.scan_queue_insert(), msg)


def test_scan_queue_request_callback(scan_guard_mock):
    sg = scan_guard_mock
    msg = messages.ScanQueueMessage(
        scan_type="fermat_scan",
        parameter={"args": {"samx": (-5, 5), "samy": (-5, 5)}, "kwargs": {"step": 3}},
        queue="primary",
    )
    msg_obj = MessageObject(MessageEndpoints.scan_queue_request(), msg)
    with mock.patch.object(sg, "_handle_scan_request") as handle:
        sg._scan_queue_request_callback(msg_obj, sg)
        handle.assert_called_once_with(msg)


def test_scan_queue_modification_request_callback(scan_guard_mock):
    sg = scan_guard_mock
    msg = messages.ScanQueueModificationMessage(
        scanID="scanID", action="abort", parameter={}, metadata={"RID": "RID"}
    )
    msg_obj = MessageObject(MessageEndpoints.scan_queue_modification(), msg)
    with mock.patch.object(sg, "_handle_scan_modification_request") as handle:
        sg._scan_queue_modification_request_callback(msg_obj, sg)
        handle.assert_called_once_with(msg)


# def test_scan_queue_modification_request_callback_wrong_msg(scan_guard_mock):
#    sg = scan_guard_mock
#    msg = messages.ScanQueueMessage(
#        scan_type="fermat_scan",
#        parameter={"args": {"samx": (-5, 5), "samy": (-5, 5)}, "kwargs": {"step": 3}},
#        queue="primary",
#    )
#    msg_obj = MessageObject(MessageEndpoints.scan_queue_modification(), msg)
#    with mock.patch.object(sg, "_handle_scan_modification_request") as handle:
#        sg._scan_queue_modification_request_callback(msg_obj, sg)
#        handle.assert_not_called()


def test_send_scan_request_response(scan_guard_mock):
    sg = scan_guard_mock
    with mock.patch.object(sg.device_manager.connector, "send") as send:
        sg._send_scan_request_response(ScanStatus(), {"RID": "RID"})
        send.assert_called_once_with(
            MessageEndpoints.scan_queue_request_response(),
            messages.RequestResponseMessage(accepted=True, message="", metadata={"RID": "RID"}),
        )


def test_handle_scan_request(scan_guard_mock):
    sg = scan_guard_mock
    msg = messages.ScanQueueMessage(
        scan_type="fermat_scan",
        parameter={"args": {"samx": (-5, 5), "samy": (-5, 5)}, "kwargs": {"step": 3}},
        queue="primary",
    )
    with mock.patch.object(sg, "_is_valid_scan_request") as valid:
        with mock.patch.object(sg, "_append_to_scan_queue") as append:
            valid.return_value = ScanStatus(accepted=True, message="")
            sg._handle_scan_request(msg)
            append.assert_called_once_with(msg)


def test_handle_scan_request_rejected(scan_guard_mock):
    sg = scan_guard_mock
    msg = messages.ScanQueueMessage(
        scan_type="fermat_scan",
        parameter={"args": {"samx": (-5, 5), "samy": (-5, 5)}, "kwargs": {"step": 3}},
        queue="primary",
    )
    with mock.patch.object(sg, "_is_valid_scan_request") as valid:
        with mock.patch.object(sg, "_append_to_scan_queue") as append:
            valid.return_value = ScanStatus(accepted=False, message="")
            sg._handle_scan_request(msg)
            append.assert_not_called()


def test_is_valid_scan_request_returns_scan_status_on_error(scan_guard_mock):
    sg = scan_guard_mock
    msg = messages.ScanQueueMessage(
        scan_type="fermat_scan",
        parameter={"args": {"samx": (-5, 5), "samy": (-5, 5)}, "kwargs": {"step": 3}},
        queue="primary",
    )
    with mock.patch.object(sg, "_check_valid_scan") as valid:
        valid.side_effect = Exception("Test exception")
        status = sg._is_valid_scan_request(msg)
        assert status.accepted == False
        assert "Test exception" in status.message


def test_check_valid_request_raises_for_empty_request(scan_guard_mock):
    sg = scan_guard_mock
    with pytest.raises(ScanRejection) as scan_rejection:
        sg._check_valid_request(None)
    assert "Invalid request." in scan_rejection.value.args