mirror of
https://github.com/bec-project/bec_widgets.git
synced 2026-10-29 00:11:28 +01:00
155 lines
6.1 KiB
Python
155 lines
6.1 KiB
Python
import threading
|
|
from typing import Callable
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import fakeredis
|
|
import pytest
|
|
from bec_lib.messages import (
|
|
ProcedureExecutionMessage,
|
|
ProcedureRequestMessage,
|
|
ProcedureWorkerStatus,
|
|
)
|
|
from bec_lib.procedures.helper import BackendProcedureHelper
|
|
from bec_lib.redis_connector import RedisConnector
|
|
from bec_server.procedures.manager import ProcedureManager
|
|
from bec_server.procedures.procedure_registry import register
|
|
from bec_server.procedures.worker_base import ProcedureWorker
|
|
|
|
import bec_widgets.tests.utils as test_utils
|
|
from bec_widgets.widgets.control.procedure_control.procedure_control import (
|
|
ProcedureControl,
|
|
QueueItem,
|
|
)
|
|
|
|
RELEASE = threading.Event()
|
|
|
|
|
|
class MockWorker(ProcedureWorker):
|
|
"""In-process stand-in for the out-of-process worker: same queue/active-set protocol."""
|
|
|
|
def _kill_process(self): ...
|
|
|
|
def _setup_execution_environment(self):
|
|
self._abort_current = threading.Event()
|
|
|
|
def _ending_or_ended(self) -> bool:
|
|
return self._aborted.is_set()
|
|
|
|
def logs(self) -> list[str]:
|
|
return []
|
|
|
|
def abort_execution(self, execution_id: str):
|
|
self._abort_current.set()
|
|
|
|
def work(self):
|
|
while (
|
|
item := self._conn.blocking_list_pop_to_set_add(
|
|
self.key, self._active_procs_endpoint, timeout_s=0.1
|
|
)
|
|
) is not None:
|
|
self.status = ProcedureWorkerStatus.RUNNING
|
|
self._helper.status_update(item.execution_id, "Started")
|
|
self._helper.notify_watchers(self._queue, queue_type="execution")
|
|
# hold the job until the test releases it or aborts it
|
|
while not (RELEASE.is_set() or self._abort_current.is_set()):
|
|
RELEASE.wait(0.02)
|
|
aborted = self._abort_current.is_set()
|
|
self._abort_current.clear()
|
|
self._helper.status_update(item.execution_id, "Aborted" if aborted else "Finished")
|
|
self._helper.remove_from_active.by_exec_id(item.execution_id)
|
|
self.status = ProcedureWorkerStatus.IDLE
|
|
self.status = ProcedureWorkerStatus.FINISHED
|
|
|
|
|
|
@pytest.fixture(scope="module", autouse=True)
|
|
def register_test_proc():
|
|
register("test", lambda: None)
|
|
|
|
|
|
@pytest.fixture
|
|
def proc_ctrl_w_helper(qtbot, mocked_client: MagicMock):
|
|
# Earlier tests shut the shared fake connector down and its pub/sub listener does not restart,
|
|
# but these tests need live notifications: give the client a fresh one.
|
|
test_utils._REDIS_CONN = None
|
|
mocked_client.connector = test_utils.global_mock_qt_redis_connector()
|
|
for key in mocked_client.connector.keys("*procedure*"):
|
|
mocked_client.connector._managed_connection.delete(key)
|
|
proc_ctrl = ProcedureControl(client=mocked_client)
|
|
qtbot.addWidget(proc_ctrl)
|
|
# the manager gets its own connection to the same fake server: its shutdown() closes it
|
|
fake_conn = lambda *_, **__: RedisConnector("localhost:1", redis_cls=fakeredis.FakeRedis)
|
|
with (
|
|
patch("bec_server.procedures.manager.RedisConnector", fake_conn),
|
|
patch("bec_server.procedures.worker_base.RedisConnector", fake_conn),
|
|
):
|
|
RELEASE.clear()
|
|
manager = ProcedureManager(MagicMock(), MockWorker)
|
|
yield proc_ctrl, BackendProcedureHelper(proc_ctrl.client.connector)
|
|
RELEASE.set()
|
|
manager.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def req_msg():
|
|
return ProcedureRequestMessage(identifier="test")
|
|
|
|
|
|
@pytest.fixture
|
|
def exec_msg():
|
|
return lambda id: ProcedureExecutionMessage(identifier="test", queue="test", execution_id=id)
|
|
|
|
|
|
def test_add_proc(qtbot, proc_ctrl_w_helper: tuple[ProcedureControl, BackendProcedureHelper]):
|
|
proc_ctrl, helper = proc_ctrl_w_helper
|
|
assert proc_ctrl._active_queues.childCount() == 0
|
|
helper.request.procedure("test")
|
|
qtbot.waitUntil(lambda: proc_ctrl._active_queues.childCount() != 0, timeout=5000)
|
|
queue: QueueItem = proc_ctrl._active_queues.child(0)
|
|
assert queue.child(0).text(1) == "RUNNING"
|
|
RELEASE.set()
|
|
qtbot.waitUntil(lambda: proc_ctrl._active_queues.childCount() == 0, timeout=5000)
|
|
|
|
|
|
def test_abort(qtbot, proc_ctrl_w_helper: tuple[ProcedureControl, BackendProcedureHelper]):
|
|
proc_ctrl, helper = proc_ctrl_w_helper
|
|
assert proc_ctrl._active_queues.childCount() == 0
|
|
helper.request.procedure("test")
|
|
qtbot.waitUntil(lambda: proc_ctrl._active_queues.childCount() != 0, timeout=5000)
|
|
|
|
assert proc_ctrl._unhandled_queues.childCount() == 0
|
|
queue: QueueItem = proc_ctrl._active_queues.child(0)
|
|
queue.child(0).actions_widget.layout().itemAt(0).widget().click()
|
|
qtbot.waitUntil(lambda: proc_ctrl._active_queues.childCount() == 0, timeout=5000)
|
|
qtbot.waitUntil(lambda: proc_ctrl._unhandled_queues.childCount() != 0, timeout=5000)
|
|
|
|
|
|
def test_delete(
|
|
qtbot,
|
|
proc_ctrl_w_helper: tuple[ProcedureControl, BackendProcedureHelper],
|
|
exec_msg: Callable[[str], ProcedureExecutionMessage],
|
|
):
|
|
proc_ctrl, helper = proc_ctrl_w_helper
|
|
[helper.push.unhandled("test", exec_msg("abcd")) for _ in range(3)]
|
|
qtbot.waitUntil(lambda: proc_ctrl._unhandled_queues.childCount() != 0, timeout=5000)
|
|
queue: QueueItem = proc_ctrl._unhandled_queues.child(0)
|
|
assert queue.childCount() == 3
|
|
queue.actions_widget.layout().itemAt(0).widget().click()
|
|
qtbot.waitUntil(lambda: helper.get.unhandled_queue("test") == [], timeout=5000)
|
|
qtbot.waitUntil(lambda: proc_ctrl._unhandled_queues.childCount() == 0, timeout=5000)
|
|
|
|
|
|
def test_resubmit(
|
|
qtbot,
|
|
proc_ctrl_w_helper: tuple[ProcedureControl, BackendProcedureHelper],
|
|
exec_msg: Callable[[str], ProcedureExecutionMessage],
|
|
):
|
|
proc_ctrl, helper = proc_ctrl_w_helper
|
|
[helper.push.unhandled("test", exec_msg(f"abcd{i}")) for i in range(3)]
|
|
qtbot.waitUntil(lambda: proc_ctrl._unhandled_queues.childCount() != 0, timeout=5000)
|
|
queue: QueueItem = proc_ctrl._unhandled_queues.child(0)
|
|
assert queue.childCount() == 3
|
|
assert proc_ctrl._active_queues.childCount() == 0
|
|
queue.child(0).actions_widget.layout().itemAt(1).widget().click()
|
|
qtbot.waitUntil(lambda: len(helper.get.unhandled_queue("test")) == 2, timeout=5000)
|
|
qtbot.waitUntil(lambda: proc_ctrl._active_queues.childCount() != 0, timeout=5000)
|