diff --git a/bec_widgets/cli/rpc/rpc_base.py b/bec_widgets/cli/rpc/rpc_base.py index a5061ae4..6a4d02e7 100644 --- a/bec_widgets/cli/rpc/rpc_base.py +++ b/bec_widgets/cli/rpc/rpc_base.py @@ -289,15 +289,29 @@ class RPCBase: if wait_for_rpc_response: try: - finished = self._msg_wait_event.wait(timeout) - if not finished: - logger.error( - "GUI RPC response timeout " - f"request_id={request_id} method={method} receiver={receiver} " - f"target_gui_id={target_gui_id} object_name={self.object_name} " - f"timeout={timeout}" - ) - raise RPCResponseTimeoutError(request_id, timeout) + start = time.monotonic() + while True: + if timeout is None: + wait_timeout = 0.5 + else: + remaining = timeout - (time.monotonic() - start) + wait_timeout = max(0.0, min(0.5, remaining)) + + finished = self._msg_wait_event.wait(wait_timeout) + if finished: + break + self._fetch_rpc_response(request_id) + if self._msg_wait_event.is_set(): + break + if timeout is not None and (time.monotonic() - start) > timeout: + logger.error( + "GUI RPC response timeout " + f"request_id={request_id} method={method} receiver={receiver} " + f"target_gui_id={target_gui_id} object_name={self.object_name} " + f"timeout={timeout}" + ) + raise RPCResponseTimeoutError(request_id, timeout) + finally: self._msg_wait_event.clear() self._client.connector.unregister( @@ -326,6 +340,13 @@ class RPCBase: self._rpc_response = msg self._msg_wait_event.set() + def _fetch_rpc_response(self, request_id: str) -> None: + msg = self._client.connector.get(MessageEndpoints.gui_instruction_response(request_id)) + if msg is None: + return + self._rpc_response = cast(messages.RequestResponseMessage, msg) + self._msg_wait_event.set() + def _create_widget_from_msg_result(self, msg_result): if msg_result is None: return None diff --git a/tests/unit_tests/test_rpc_base.py b/tests/unit_tests/test_rpc_base.py index ac668654..8cd7a28d 100644 --- a/tests/unit_tests/test_rpc_base.py +++ b/tests/unit_tests/test_rpc_base.py @@ -1,6 +1,7 @@ from unittest.mock import MagicMock import pytest +from bec_lib import messages from bec_lib.device import DeviceBaseWithConfig, Signal from bec_widgets.cli.rpc import rpc_base as rpc_base_module @@ -59,6 +60,7 @@ def test_run_rpc_logs_response_timeout(monkeypatch): rpc = RPCBase(gui_id="progress_widget", object_name="progressbar") rpc._rpc_timeout = 0 rpc._client = MagicMock() + rpc._client.connector.get.return_value = None info_mock = MagicMock() error_mock = MagicMock() @@ -83,3 +85,55 @@ def test_run_rpc_logs_response_timeout(monkeypatch): assert "target_gui_id=progress_widget" in error_message assert "object_name=progressbar" in error_message assert "timeout=0" in error_message + + +def test_run_rpc_waits_indefinitely_when_timeout_is_none(monkeypatch): + rpc = RPCBase(gui_id="progress_widget", object_name="progressbar") + rpc._client = MagicMock() + rpc._create_widget_from_msg_result = MagicMock(return_value="done") + + response = messages.RequestResponseMessage(accepted=True, message={"result": "rpc-result"}) + + wait_calls = {"count": 0} + + def wait_side_effect(timeout): + wait_calls["count"] += 1 + if wait_calls["count"] == 1: + return False + rpc._rpc_response = response + rpc._msg_wait_event.set() + return True + + monkeypatch.setattr(rpc._msg_wait_event, "wait", wait_side_effect) + fetch_mock = MagicMock() + monkeypatch.setattr(rpc, "_fetch_rpc_response", fetch_mock) + + result = rpc._run_rpc("set_value", timeout=None) + + assert result == "done" + assert fetch_mock.call_count == 1 + assert wait_calls["count"] == 2 + assert rpc._create_widget_from_msg_result.call_args.args == ("rpc-result",) + + +def test_run_rpc_fetches_response_before_timeout(monkeypatch): + rpc = RPCBase(gui_id="progress_widget", object_name="progressbar") + rpc._client = MagicMock() + rpc._create_widget_from_msg_result = MagicMock(return_value="done") + + response = messages.RequestResponseMessage(accepted=True, message={"result": "rpc-result"}) + + monkeypatch.setattr(rpc._msg_wait_event, "wait", MagicMock(return_value=False)) + + def fetch_side_effect(_request_id): + rpc._rpc_response = response + rpc._msg_wait_event.set() + + fetch_mock = MagicMock(side_effect=fetch_side_effect) + monkeypatch.setattr(rpc, "_fetch_rpc_response", fetch_mock) + + result = rpc._run_rpc("set_value", timeout=0) + + assert result == "done" + fetch_mock.assert_called_once() + rpc._create_widget_from_msg_result.assert_called_once_with("rpc-result")