bec/bec_lib/tests/test_scan_manager.py

from unittest import mock

import pytest

from bec_lib import messages
from bec_lib.endpoints import MessageEndpoints
from bec_lib.scan_manager import ScanManager


@pytest.fixture
def scan_manager():
    connector = mock.MagicMock()
    manager = ScanManager(connector=connector)
    yield manager
    manager.shutdown()


def test_scan_manager_next_scan_number(scan_manager):
    scan_manager.connector.get.return_value = messages.VariableMessage(value=3)
    assert scan_manager.next_scan_number == 3


def test_scan_manager_next_scan_number_failed(scan_manager):
    scan_manager.connector.get.return_value = None
    assert scan_manager.next_scan_number == -1


def test_scan_manager_next_scan_number_with_int(scan_manager):
    scan_manager.connector.get.return_value = 3
    assert scan_manager.next_scan_number == 3


def test_scan_manager_next_scan_number_setter(scan_manager):
    scan_manager.next_scan_number = 3
    scan_manager.connector.set.assert_called_once_with(
        MessageEndpoints.scan_number(), messages.VariableMessage(value=3)
    )


def test_scan_manager_next_dataset_number(scan_manager):
    scan_manager.connector.get.return_value = messages.VariableMessage(value=3)
    assert scan_manager.next_dataset_number == 3


def test_scan_manager_next_dataset_number_failed(scan_manager):
    scan_manager.connector.get.return_value = None
    assert scan_manager.next_dataset_number == -1


def test_scan_manager_next_dataset_number_with_int(scan_manager):
    scan_manager.connector.get.return_value = 3
    assert scan_manager.next_dataset_number == 3


def test_scan_manager_next_dataset_number_setter(scan_manager):
    scan_manager.next_dataset_number = 3
    scan_manager.connector.set.assert_called_once_with(
        MessageEndpoints.dataset_number(), messages.VariableMessage(value=3)
    )


def test_scan_manager_request_scan_abortion(scan_manager):
    scan_manager.request_scan_abortion("scan_id")
    scan_manager.connector.send.assert_called_once_with(
        MessageEndpoints.scan_queue_modification_request(),
        messages.ScanQueueModificationMessage(scan_id="scan_id", action="abort", parameter={}),
    )


@pytest.mark.parametrize("scan_id", [None, "scan_id", ["scan_id"], [None]])
def test_scan_manager_request_scan_abortion_scan_id(scan_manager, scan_id):

    class ScanStorage:
        current_scan_info = {"scan_id": scan_id}

        @property
        def current_scan_id(self):
            return self.current_scan_info["scan_id"]

    scan_manager.scan_storage = ScanStorage()
    scan_manager.request_scan_abortion()
    scan_manager.connector.send.assert_called_once_with(
        MessageEndpoints.scan_queue_modification_request(),
        messages.ScanQueueModificationMessage(scan_id=scan_id, action="abort", parameter={}),
    )


def test_scan_manager_request_scan_halt(scan_manager):
    scan_manager.request_scan_halt("scan_id")
    scan_manager.connector.send.assert_called_once_with(
        MessageEndpoints.scan_queue_modification_request(),
        messages.ScanQueueModificationMessage(scan_id="scan_id", action="halt", parameter={}),
    )


@pytest.mark.parametrize("scan_id", [None, "scan_id", ["scan_id"], [None]])
def test_scan_manager_request_scan_halt_scan_id(scan_manager, scan_id):

    class ScanStorage:
        current_scan_info = {"scan_id": scan_id}

        @property
        def current_scan_id(self):
            return self.current_scan_info["scan_id"]

    scan_manager.scan_storage = ScanStorage()
    scan_manager.request_scan_halt()
    scan_manager.connector.send.assert_called_once_with(
        MessageEndpoints.scan_queue_modification_request(),
        messages.ScanQueueModificationMessage(scan_id=scan_id, action="halt", parameter={}),
    )


def test_scan_manager_request_scan_continuation(scan_manager):
    scan_manager.request_scan_continuation("scan_id")
    scan_manager.connector.send.assert_called_once_with(
        MessageEndpoints.scan_queue_modification_request(),
        messages.ScanQueueModificationMessage(scan_id="scan_id", action="continue", parameter={}),
    )


@pytest.mark.parametrize("scan_id", [None, "scan_id", ["scan_id"], [None]])
def test_scan_manager_request_scan_continuation_scan_id(scan_manager, scan_id):

    class ScanStorage:
        current_scan_info = {"scan_id": scan_id}

        @property
        def current_scan_id(self):
            return self.current_scan_info["scan_id"]

    scan_manager.scan_storage = ScanStorage()
    scan_manager.request_scan_continuation()
    scan_manager.connector.send.assert_called_once_with(
        MessageEndpoints.scan_queue_modification_request(),
        messages.ScanQueueModificationMessage(scan_id=scan_id, action="continue", parameter={}),
    )