bec-ipython-client 3.35.6__py3-none-any.whl → 3.84.0__py3-none-any.whl

tests/client_tests/test_move_callback.py

@@ -1,20 +1,23 @@
 import collections
+import time
 from unittest import mock
 
 import pytest
 
 from bec_ipython_client.callbacks.move_device import (
     LiveUpdatesReadbackProgressbar,
-    ReadbackDataMixin,
+    ReadbackDataHandler,
 )
 from bec_lib import messages
 from bec_lib.endpoints import MessageEndpoints
 
 
 @pytest.fixture
-def readback_data_mixin(bec_client_mock):
-    with mock.patch.object(bec_client_mock.device_manager, "connector"):
-        yield ReadbackDataMixin(bec_client_mock.device_manager, ["samx", "samy"])
+def readback_data_handler(bec_client_mock, connected_connector):
+    with mock.patch.object(bec_client_mock.device_manager, "connector", connected_connector):
+        yield ReadbackDataHandler(
+            bec_client_mock.device_manager, ["samx", "samy"], request_id="something"
+        )
 
 
 def test_move_callback(bec_client_mock):
@@ -33,30 +36,25 @@ def test_move_callback(bec_client_mock):
         return readback[0]
 
     req_done = collections.deque()
-    msg_acc = messages.DeviceReqStatusMessage(
-        device="samx", success=True, metadata={"RID": "something"}
-    )
-    req_done.extend([[None], [None], [None], [msg_acc]])
+    req_done.extend([{"samx": (False, False)}, {"samx": (False, False)}, {"samx": (True, True)}])
 
     def mock_req_msg(*args):
         if len(req_done) > 1:
             return req_done.popleft()
         return req_done[0]
 
-    with mock.patch("bec_ipython_client.callbacks.move_device.check_alarms") as check_alarms_mock:
-        with mock.patch.object(ReadbackDataMixin, "wait_for_RID"):
-            with mock.patch.object(LiveUpdatesReadbackProgressbar, "wait_for_request_acceptance"):
+    with mock.patch("bec_ipython_client.callbacks.move_device.check_alarms"):
+        with mock.patch.object(LiveUpdatesReadbackProgressbar, "wait_for_request_acceptance"):
+            with mock.patch.object(
+                LiveUpdatesReadbackProgressbar, "_print_client_msgs_asap"
+            ) as mock_client_msgs:
                 with mock.patch.object(
-                    LiveUpdatesReadbackProgressbar, "_print_client_msgs_asap"
-                ) as mock_client_msgs:
-                    with mock.patch.object(
-                        LiveUpdatesReadbackProgressbar, "_print_client_msgs_all"
-                    ) as mock_client_msgs_all:
-                        with mock.patch.object(
-                            ReadbackDataMixin, "get_device_values", mock_readback
-                        ):
+                    LiveUpdatesReadbackProgressbar, "_print_client_msgs_all"
+                ) as mock_client_msgs_all:
+                    with mock.patch.object(ReadbackDataHandler, "get_device_values", mock_readback):
+                        with mock.patch.object(ReadbackDataHandler, "device_states", mock_req_msg):
                             with mock.patch.object(
-                                ReadbackDataMixin, "get_request_done_msgs", mock_req_msg
+                                ReadbackDataHandler, "done", side_effect=[False, False, True]
                             ):
                                 LiveUpdatesReadbackProgressbar(bec=client, request=request).run()
                                 assert mock_client_msgs.called is True
@@ -82,28 +80,21 @@ def test_move_callback_with_report_instruction(bec_client_mock):
         return readback[0]
 
     req_done = collections.deque()
-    msg_acc = messages.DeviceReqStatusMessage(
-        device="samx", success=True, metadata={"RID": "something"}
-    )
-    req_done.extend([[None], [None], [None], [msg_acc]])
+    req_done.extend([{"samx": (False, False)}, {"samx": (False, False)}, {"samx": (True, True)}])
 
     def mock_req_msg(*args):
         if len(req_done) > 1:
             return req_done.popleft()
         return req_done[0]
 
-    with mock.patch("bec_ipython_client.callbacks.move_device.check_alarms") as check_alarms_mock:
-        with mock.patch.object(ReadbackDataMixin, "wait_for_RID"):
-            with mock.patch.object(LiveUpdatesReadbackProgressbar, "wait_for_request_acceptance"):
-                with mock.patch.object(LiveUpdatesReadbackProgressbar, "_print_client_msgs_asap"):
-                    with mock.patch.object(
-                        LiveUpdatesReadbackProgressbar, "_print_client_msgs_all"
-                    ):
-                        with mock.patch.object(
-                            ReadbackDataMixin, "get_device_values", mock_readback
-                        ):
+    with mock.patch("bec_ipython_client.callbacks.move_device.check_alarms"):
+        with mock.patch.object(LiveUpdatesReadbackProgressbar, "wait_for_request_acceptance"):
+            with mock.patch.object(LiveUpdatesReadbackProgressbar, "_print_client_msgs_asap"):
+                with mock.patch.object(LiveUpdatesReadbackProgressbar, "_print_client_msgs_all"):
+                    with mock.patch.object(ReadbackDataHandler, "get_device_values", mock_readback):
+                        with mock.patch.object(ReadbackDataHandler, "device_states", mock_req_msg):
                             with mock.patch.object(
-                                ReadbackDataMixin, "get_request_done_msgs", mock_req_msg
+                                ReadbackDataHandler, "done", side_effect=[False, False, False, True]
                             ):
                                 LiveUpdatesReadbackProgressbar(
                                     bec=client,
@@ -112,70 +103,121 @@ def test_move_callback_with_report_instruction(bec_client_mock):
                                 ).run()
 
 
-def test_readback_data_mixin(readback_data_mixin):
-    readback_data_mixin.device_manager.connector.get.side_effect = [
-        messages.DeviceMessage(
+def test_readback_data_handler(readback_data_handler):
+    readback_data_handler.data = {
+        "samx": messages.DeviceMessage(
             signals={"samx": {"value": 10}, "samx_setpoint": {"value": 20}},
             metadata={"device": "samx"},
         ),
-        messages.DeviceMessage(
+        "samy": messages.DeviceMessage(
             signals={"samy": {"value": 10}, "samy_setpoint": {"value": 20}},
             metadata={"device": "samy"},
         ),
-    ]
-    res = readback_data_mixin.get_device_values()
+    }
+
+    res = readback_data_handler.get_device_values()
     assert res == [10, 10]
 
 
-def test_readback_data_mixin_multiple_hints(readback_data_mixin):
-    readback_data_mixin.device_manager.devices.samx._info["hints"]["fields"] = [
+def test_readback_data_handler_multiple_hints(readback_data_handler):
+    readback_data_handler.device_manager.devices.samx._info["hints"]["fields"] = [
         "samx_setpoint",
         "samx",
     ]
-    readback_data_mixin.device_manager.connector.get.side_effect = [
-        messages.DeviceMessage(
+    readback_data_handler.data = {
+        "samx": messages.DeviceMessage(
             signals={"samx": {"value": 10}, "samx_setpoint": {"value": 20}},
             metadata={"device": "samx"},
         ),
-        messages.DeviceMessage(
+        "samy": messages.DeviceMessage(
             signals={"samy": {"value": 10}, "samy_setpoint": {"value": 20}},
             metadata={"device": "samy"},
         ),
-    ]
-    res = readback_data_mixin.get_device_values()
+    }
+    res = readback_data_handler.get_device_values()
     assert res == [20, 10]
 
 
-def test_readback_data_mixin_multiple_no_hints(readback_data_mixin):
-    readback_data_mixin.device_manager.devices.samx._info["hints"]["fields"] = []
-    readback_data_mixin.device_manager.connector.get.side_effect = [
-        messages.DeviceMessage(
+def test_readback_data_handler_multiple_no_hints(readback_data_handler):
+    readback_data_handler.device_manager.devices.samx._info["hints"]["fields"] = []
+    readback_data_handler.data = {
+        "samx": messages.DeviceMessage(
             signals={"samx": {"value": 10}, "samx_setpoint": {"value": 20}},
             metadata={"device": "samx"},
         ),
-        messages.DeviceMessage(
+        "samy": messages.DeviceMessage(
             signals={"samy": {"value": 10}, "samy_setpoint": {"value": 20}},
             metadata={"device": "samy"},
         ),
-    ]
-    res = readback_data_mixin.get_device_values()
+    }
+    res = readback_data_handler.get_device_values()
     assert res == [10, 10]
 
 
-def test_get_request_done_msgs(readback_data_mixin):
-    res = readback_data_mixin.get_request_done_msgs()
-    readback_data_mixin.device_manager.connector.pipeline.assert_called_once()
-    assert (
-        mock.call(
-            MessageEndpoints.device_req_status("samx"),
-            readback_data_mixin.device_manager.connector.pipeline.return_value,
-        )
-        in readback_data_mixin.device_manager.connector.get.call_args_list
+def test_readback_data_handler_init(readback_data_handler):
+    """
+    Test that the ReadbackDataHandler is initialized correctly.
+    """
+
+    # Initial state
+    assert readback_data_handler._devices_done_state == {
+        "samx": (False, False),
+        "samy": (False, False),
+    }
+    assert readback_data_handler._devices_received == {"samx": False, "samy": False}
+    assert readback_data_handler.data == {}
+
+
+def test_readback_data_handler_readback_callbacks(readback_data_handler):
+    """
+    Test that the readback callback properly updates the readback data.
+    """
+
+    # Submit readback for samx
+    msg = messages.DeviceMessage(
+        signals={"samx": {"value": 15}}, metadata={"device": "samx", "RID": "something"}
     )
-    assert (
-        mock.call(
-            MessageEndpoints.device_req_status("samy"),
-            readback_data_mixin.device_manager.connector.pipeline.return_value,
-        )
-        in readback_data_mixin.device_manager.connector.get.call_args_list
+    readback_data_handler.connector.set_and_publish(MessageEndpoints.device_readback("samx"), msg)
+
+    msg_old = messages.DeviceMessage(
+        signals={"samy": {"value": 10}}, metadata={"device": "samx", "RID": "something_else"}
+    )
+    readback_data_handler.connector.set_and_publish(
+        MessageEndpoints.device_readback("samy"), msg_old
+    )
+    while (
+        not readback_data_handler._devices_received["samx"]
+        or "samx" not in readback_data_handler.data
+    ):
+        time.sleep(0.01)
+    assert readback_data_handler.data["samx"].signals["samx"]["value"] == 15
+    dev_data = readback_data_handler.get_device_values()
+    assert dev_data[0] == 15
+    assert dev_data[1] == 10  # samy remains unchanged
+
+
+def test_readback_data_handler_request_done_callbacks(readback_data_handler):
+    """
+    Test that the request done callback properly updates the device done state.
+    """
+
+    # Submit request done for samx
+    msg = messages.DeviceReqStatusMessage(device="samx", success=True, request_id="something")
+    readback_data_handler.connector.xadd(
+        MessageEndpoints.device_req_status("something"), {"data": msg}
+    )
+    while not readback_data_handler._devices_done_state["samx"][0]:
+        time.sleep(0.01)
+    assert readback_data_handler._devices_done_state["samx"] == (True, True)
+
+    assert readback_data_handler.done() is False
+
+    # Submit request done for samy
+    msg = messages.DeviceReqStatusMessage(device="samy", success=False, request_id="something")
+    readback_data_handler.connector.xadd(
+        MessageEndpoints.device_req_status("something"), {"data": msg}
     )
+    while not readback_data_handler._devices_done_state["samy"][0]:
+        time.sleep(0.01)
+    assert readback_data_handler._devices_done_state["samy"] == (True, False)
+    assert readback_data_handler.done() is True

tests/end-2-end/_ensure_requirements_container.py

@@ -0,0 +1,21 @@
+from time import sleep
+
+from bec_server.scan_server.procedures.constants import PROCEDURE, ProcedureWorkerError
+from bec_server.scan_server.procedures.container_utils import PodmanCliUtils
+
+image_name = (
+    f"ghcr.io/bec-project/{PROCEDURE.CONTAINER.REQUIREMENTS_IMAGE_NAME}:v{PROCEDURE.BEC_VERSION}"
+)
+podman = PodmanCliUtils()
+
+for i in range(1, 6):
+    try:
+        output = podman._run_and_capture_error("podman", "pull", image_name)
+        print("successfully pulled requirements image for current version")
+        exit(0)
+    except ProcedureWorkerError as e:
+        print(e)
+        print("retrying in 5 minutes...")
+        sleep(5 * 60)
+print(f"No more retries. Check if {image_name} actually exists!")
+exit(1)

tests/end-2-end/test_procedures_e2e.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from importlib.metadata import version
+from typing import TYPE_CHECKING, Callable, Generator
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from bec_ipython_client.main import BECIPythonClient
+from bec_lib import messages
+from bec_lib.endpoints import MessageEndpoints
+from bec_lib.logger import bec_logger
+from bec_server.scan_server.procedures.constants import _CONTAINER, _WORKER
+from bec_server.scan_server.procedures.container_utils import get_backend
+from bec_server.scan_server.procedures.container_worker import ContainerProcedureWorker
+from bec_server.scan_server.procedures.manager import ProcedureManager
+
+if TYPE_CHECKING:
+    from pytest_bec_e2e.plugin import LogTestTool
+
+logger = bec_logger.logger
+
+# pylint: disable=protected-access
+
+# Random order disabled for this module so that the test for building the worker container runs first
+# and we can use lower timeouts for the remaining tests
+pytestmark = pytest.mark.random_order(disabled=True)
+
+
+@dataclass(frozen=True)
+class PATCHED_CONSTANTS:
+    WORKER = _WORKER()
+    CONTAINER = _CONTAINER()
+    MANAGER_SHUTDOWN_TIMEOUT_S = 2
+    BEC_VERSION = version("bec_lib")
+    REDIS_HOST = "localhost"
+
+
+@pytest.fixture
+def client_logtool_and_manager(
+    bec_ipython_client_fixture_with_logtool: tuple[BECIPythonClient, "LogTestTool"],
+) -> Generator[tuple[BECIPythonClient, "LogTestTool", ProcedureManager], None, None]:
+    client, logtool = bec_ipython_client_fixture_with_logtool
+    server = MagicMock()
+    server.bootstrap_server = f"{client.connector.host}:{client.connector.port}"
+    manager = ProcedureManager(server, ContainerProcedureWorker)
+    yield client, logtool, manager
+    manager.shutdown()
+
+
+def _wait_while(cond: Callable[[], bool], timeout_s):
+    start = time.monotonic()
+    while cond():
+        if (time.monotonic() - start) > timeout_s:
+            raise TimeoutError()
+        time.sleep(0.01)
+
+
+@pytest.mark.timeout(100)
+def test_building_worker_image():
+    podman_utils = get_backend()
+    build = podman_utils.build_worker_image()
+    assert len(build._command_output.splitlines()[-1]) == 64  # type: ignore
+    assert podman_utils.image_exists(f"bec_procedure_worker:v{version('bec_lib')}")
+
+
+@pytest.mark.timeout(100)
+@patch("bec_server.scan_server.procedures.manager.procedure_registry.is_registered", lambda _: True)
+def test_procedure_runner_spawns_worker(
+    client_logtool_and_manager: tuple[BECIPythonClient, "LogTestTool", ProcedureManager],
+):
+    client, _, manager = client_logtool_and_manager
+    assert manager._active_workers == {}
+    endpoint = MessageEndpoints.procedure_request()
+    msg = messages.ProcedureRequestMessage(
+        identifier="sleep", args_kwargs=((), {"time_s": 2}), queue="test"
+    )
+
+    logs = []
+
+    def cb(worker: ContainerProcedureWorker):
+        nonlocal logs
+        logs = worker._backend.logs(worker._container_id)
+
+    manager.add_callback("test", cb)
+    client.connector.xadd(topic=endpoint, msg_dict=msg.model_dump())
+
+    _wait_while(lambda: manager._active_workers == {}, 5)
+    _wait_while(lambda: manager._active_workers != {}, 20)
+
+    assert logs != []
+
+
+@pytest.mark.timeout(100)
+@patch("bec_server.scan_server.procedures.manager.procedure_registry.is_registered", lambda _: True)
+@patch("bec_server.scan_server.procedures.container_worker.PROCEDURE", PATCHED_CONSTANTS())
+def test_happy_path_container_procedure_runner(
+    client_logtool_and_manager: tuple[BECIPythonClient, "LogTestTool", ProcedureManager],
+):
+    test_args = (1, 2, 3)
+    test_kwargs = {"a": "b", "c": "d"}
+    client, logtool, manager = client_logtool_and_manager
+    assert manager._active_workers == {}
+    conn = client.connector
+    endpoint = MessageEndpoints.procedure_request()
+    msg = messages.ProcedureRequestMessage(
+        identifier="log execution message args", args_kwargs=(test_args, test_kwargs)
+    )
+    conn.xadd(topic=endpoint, msg_dict=msg.model_dump())
+
+    _wait_while(lambda: manager._active_workers == {}, 5)
+    _wait_while(lambda: manager._active_workers != {}, 20)
+
+    logtool.fetch()
+    assert logtool.is_present_in_any_message("procedure accepted: True, message:")
+    assert logtool.is_present_in_any_message(
+        "ContainerWorker started container for queue primary"
+    ), f"Log content relating to procedures: {manager._logs}"
+
+    res, msg = logtool.are_present_in_order(
+        [
+            "Container worker 'primary' status update: IDLE",
+            "Container worker 'primary' status update: RUNNING",
+            "Container worker 'primary' status update: IDLE",
+            "Container worker 'primary' status update: FINISHED",
+        ]
+    )
+    assert res, f"failed on {msg}"
+
+    assert logtool.is_present_in_any_message(
+        f"Builtin procedure log_message_args_kwargs called with args: {test_args} and kwargs: {test_kwargs}"
+    )

tests/end-2-end/test_scans_e2e.py

@@ -737,12 +737,12 @@ def test_update_config(bec_ipython_client_fixture):
     bec = bec_ipython_client_fixture
     bec.metadata.update({"unit_test": "test_update_config"})
     demo_config_path = os.path.join(os.path.dirname(configs.__file__), "demo_config.yaml")
-    config = bec.config._load_config_from_file(demo_config_path)
+    config = bec.device_manager.config_helper._load_config_from_file(demo_config_path)
     config.pop("samx")
-    bec.config.send_config_request(action="set", config=config)
+    bec.device_manager.config_helper.send_config_request(action="set", config=config)
     assert "samx" not in bec.device_manager.devices
-    config = bec.config._load_config_from_file(demo_config_path)
-    bec.config.send_config_request(action="set", config=config)
+    config = bec.device_manager.config_helper._load_config_from_file(demo_config_path)
+    bec.device_manager.config_helper.send_config_request(action="set", config=config)
 
 
 @pytest.mark.timeout(100)