diff --git a/pixi.toml b/pixi.toml index 48fd7586..5270b82a 100644 --- a/pixi.toml +++ b/pixi.toml @@ -3,9 +3,10 @@ authors = ["Dmitri Gavrilov "] channels = ["conda-forge"] name = "bluesky-queueserver" platforms = ["linux-64"] -version = "0.0.24" +version = "0.0.25" [tasks] +lint = "pre-commit run --all-files" [dependencies] python = "*" diff --git a/src/bluesky_queueserver/__init__.py b/src/bluesky_queueserver/__init__.py index 9294a94e..afecfd2d 100644 --- a/src/bluesky_queueserver/__init__.py +++ b/src/bluesky_queueserver/__init__.py @@ -12,6 +12,8 @@ from .manager.output_streaming import ( # noqa: E402, F401 ReceiveConsoleOutput, ReceiveConsoleOutputAsync, + ReceiveProgressInfo, + ReceiveProgressInfoAsync, ReceiveSystemInfo, ReceiveSystemInfoAsync, ) diff --git a/src/bluesky_queueserver/manager/config.py b/src/bluesky_queueserver/manager/config.py index 002add5f..6d55b1b0 100644 --- a/src/bluesky_queueserver/manager/config.py +++ b/src/bluesky_queueserver/manager/config.py @@ -185,6 +185,8 @@ def parse_configs(config_path): "zmq_info_addr": "network/zmq_info_addr", "zmq_encoding": "network/zmq_encoding", "zmq_publish_console": "network/zmq_publish_console", + "zmq_publish_info": "network/zmq_publish_info", + "zmq_publish_progress": "network/zmq_publish_progress", "redis_addr": "network/redis_addr", "redis_name_prefix": "network/redis_name_prefix", "ignore_invalid_plans": "startup/ignore_invalid_plans", @@ -355,6 +357,18 @@ def __init__(self, *, parser, args): value_cli=self._args_existing("zmq_publish_console"), ) + self._settings["zmq_publish_info"] = self._get_param_boolean( + value_default=args.zmq_publish_info, + value_config=self._get_value_from_config("zmq_publish_info"), + value_cli=self._args_existing("zmq_publish_info"), + ) + + self._settings["zmq_publish_progress"] = self._get_param_boolean( + value_default=args.zmq_publish_progress, + value_config=self._get_value_from_config("zmq_publish_progress"), + value_cli=self._args_existing("zmq_publish_progress"), + ) + redis_addr = self._get_param( value_default=self._args.redis_addr, value_config=self._get_value_from_config("redis_addr"), diff --git a/src/bluesky_queueserver/manager/config_schemas/config_schema.yml b/src/bluesky_queueserver/manager/config_schemas/config_schema.yml index 74379ecc..e950a265 100644 --- a/src/bluesky_queueserver/manager/config_schemas/config_schema.yml +++ b/src/bluesky_queueserver/manager/config_schemas/config_schema.yml @@ -21,6 +21,10 @@ properties: type: string zmq_publish_console: type: boolean + zmq_publish_info: + type: boolean + zmq_publish_progress: + type: boolean redis_addr: type: string redis_name_prefix: diff --git a/src/bluesky_queueserver/manager/output_streaming.py b/src/bluesky_queueserver/manager/output_streaming.py index 7f420869..dc994c7c 100644 --- a/src/bluesky_queueserver/manager/output_streaming.py +++ b/src/bluesky_queueserver/manager/output_streaming.py @@ -115,8 +115,29 @@ def push_info_to_msg_queue(*, key, msg, msg_queue): msg_queue.put(msg) +def push_progress_to_msg_queue(*, msg, msg_queue): + """ + Format a progress message and put it into the message queue. The message is published + to the ``progress`` channel on the ``QS_Progress`` 0MQ topic. + + Parameters + ---------- + msg : dict + The progress payload dictionary. + msg_queue : multiprocessing.Queue + Reference to the queue used for collecting messages. + + Returns + ------- + None + """ + msg = {"channel": "progress", "time": ttime.time(), "msg": msg} + msg_queue.put(msg) + + _default_zmq_console_topic = "QS_Console" _default_zmq_info_topic = "QS_Info" +_default_zmq_progress_topic = "QS_Progress" class PublishZMQStreamOutput: @@ -135,8 +156,12 @@ class PublishZMQStreamOutput: The messages added to the queue will be automatically published to 0MQ socket. console_output_on : boolean Enable/disable printing console output to the terminal - zmq_publish_on : boolean - Enable/disable publishing console output to 0MQ socket + zmq_publish_console : boolean + Enable/disable publishing console output to 0MQ socket (``QS_Console`` topic). + zmq_publish_info : boolean + Enable/disable publishing info/status messages to 0MQ socket (``QS_Info`` topic). + zmq_publish_progress : boolean + Enable/disable publishing progress messages to 0MQ socket (``QS_Progress`` topic). zmq_publish_addr : str, None Address of 0MQ PUB socket for the publishing server. If ``None``, then the default address ``tcp://*:60625`` is used. @@ -146,6 +171,8 @@ class PublishZMQStreamOutput: Name of the 0MQ topic where the console messages are published. zmq_topic_info : str Name of the 0MQ topic where the system information messages are published. + zmq_topic_progress : str + Name of the 0MQ topic where the progress messages are published. name : str Name of the thread where the messages are published. """ @@ -155,11 +182,14 @@ def __init__( *, msg_queue, console_output_on=True, - zmq_publish_on=True, + zmq_publish_console=True, + zmq_publish_info=True, + zmq_publish_progress=True, zmq_publish_addr=None, encoding="json", zmq_topic_console=_default_zmq_console_topic, zmq_topic_info=_default_zmq_info_topic, + zmq_topic_progress=_default_zmq_progress_topic, name="RE Console Output Publisher", ): self._thread_running = False # Set True to exit the thread @@ -168,7 +198,9 @@ def __init__( self._polling_timeout = 0.1 # in sec. self._console_output_on = console_output_on - self._zmq_publish_on = zmq_publish_on + self._zmq_publish_console = zmq_publish_console + self._zmq_publish_info = zmq_publish_info + self._zmq_publish_progress = zmq_publish_progress self._encoding = process_zmq_encoding_name(encoding) @@ -177,22 +209,24 @@ def __init__( self._zmq_publish_addr = zmq_publish_addr self._zmq_topic_console = zmq_topic_console self._zmq_topic_info = zmq_topic_info + self._zmq_topic_progress = zmq_topic_progress + zmq_publish_on = zmq_publish_console or zmq_publish_info or zmq_publish_progress self._socket = None - if self._zmq_publish_on: + if zmq_publish_on: try: context = zmq.Context() self._socket = context.socket(zmq.PUB) self._socket.bind(self._zmq_publish_addr) except Exception as ex: logger.error( - "Failed to create 0MQ socket at %s. Console output will not be published. Exception: %s", + "Failed to create 0MQ socket at %s. Output will not be published. Exception: %s", self._zmq_publish_addr, ex, ) - if self._socket and self._zmq_publish_on: - logging.info("Publishing console output to 0MQ socket at %s", zmq_publish_addr) + if self._socket and zmq_publish_on: + logging.info("Publishing output to 0MQ socket at %s", zmq_publish_addr) def start(self): """ @@ -237,13 +271,15 @@ def _publish(self, payload): sys.__stdout__.write(payload["msg"]) sys.__stdout__.flush() - if self._zmq_publish_on and self._socket: - if channel == "console": + if self._socket: + if channel == "console" and self._zmq_publish_console: topic = self._zmq_topic_console - elif channel == "info": + elif channel == "info" and self._zmq_publish_info: topic = self._zmq_topic_info + elif channel == "progress" and self._zmq_publish_progress: + topic = self._zmq_topic_progress else: - logger.error("Failed to publish the message: unsupported 0MQ channel %s.") + return payload = {k: payload[k] for k in ("time", "msg")} if self._encoding == ZMQEncoding.JSON: payload_json = json.dumps(payload) @@ -427,8 +463,26 @@ def __init__( ) +class ReceiveProgressInfo(_ReceiveZMQStreamOutput): + """ + The class defaults are set to receive 0MQ messages with progress information + (RunEngine waiting/watcher updates). + """ + + def __init__( + self, *, zmq_subscribe_addr=None, encoding="json", zmq_topic=_default_zmq_progress_topic, timeout=1000 + ): + super().__init__( + zmq_subscribe_addr=zmq_subscribe_addr, + encoding=encoding, + zmq_topic=zmq_topic, + timeout=timeout, + ) + + ReceiveConsoleOutput.__doc__ += _ReceiveZMQStreamOutput.__doc__ ReceiveSystemInfo.__doc__ += _ReceiveZMQStreamOutput.__doc__ +ReceiveProgressInfo.__doc__ += _ReceiveZMQStreamOutput.__doc__ class _ReceiveZMQStreamOutputAsync: @@ -717,8 +771,26 @@ def __init__( ) +class ReceiveProgressInfoAsync(_ReceiveZMQStreamOutputAsync): + """ + The class defaults are set to receive 0MQ messages with progress information + (RunEngine waiting/watcher updates). + """ + + def __init__( + self, *, zmq_subscribe_addr=None, encoding="json", zmq_topic=_default_zmq_progress_topic, timeout=1000 + ): + super().__init__( + zmq_subscribe_addr=zmq_subscribe_addr, + encoding=encoding, + zmq_topic=zmq_topic, + timeout=timeout, + ) + + ReceiveConsoleOutputAsync.__doc__ += _ReceiveZMQStreamOutputAsync.__doc__ ReceiveSystemInfoAsync.__doc__ += _ReceiveZMQStreamOutputAsync.__doc__ +ReceiveProgressInfoAsync.__doc__ += _ReceiveZMQStreamOutputAsync.__doc__ def qserver_console_monitor_cli(): @@ -744,9 +816,10 @@ def formatter(prog): dest="zmq_info_addr", type=str, default=None, - help="The address of RE Manager socket used for publishing console output. The parameter overrides " - "the address set using QSERVER_ZMQ_INFO_ADDRESS environment variable. The default value is used " - "if the address is not set using the parameter or the environment variable. Address format: " + help="The address of RE Manager socket used for publishing console output, status info, and " + "progress updates. The parameter overrides the address set using QSERVER_ZMQ_INFO_ADDRESS " + "environment variable. The default value is used if the address is not set using the parameter " + "or the environment variable. Address format: " f"'tcp://127.0.0.1:60625' (default: {default_zmq_info_address}).", ) diff --git a/src/bluesky_queueserver/manager/plan_monitoring.py b/src/bluesky_queueserver/manager/plan_monitoring.py index f1b2f00c..8eb1b5de 100644 --- a/src/bluesky_queueserver/manager/plan_monitoring.py +++ b/src/bluesky_queueserver/manager/plan_monitoring.py @@ -1,9 +1,12 @@ import copy import logging import threading +import time as ttime from bluesky.callbacks.core import CallbackBase +from .output_streaming import push_progress_to_msg_queue + logger = logging.getLogger(__name__) @@ -205,3 +208,150 @@ def stop(self, doc): logger.info("Run was closed: %r", uid) except Exception as ex: logger.exception("RE Manager: Failed to label run as closed: %s", ex) + + +def _to_json_safe(value): + """ + Coerce a value to a JSON-serializable type. Returns ``None`` for values + that cannot be represented as a number or string. + """ + if value is None: + return None + if isinstance(value, (int, float, bool)): + return value + if isinstance(value, str): + return value + try: + return float(value) + except (TypeError, ValueError): + return str(value) + + +class WatcherStreamManager: + """ + RunEngine ``waiting_hook``-compatible class. Instead of rendering progress bars, + it serializes watcher updates and pushes them to ``msg_queue`` on the ``"progress"`` + channel (``QS_Progress`` 0MQ topic) so they are published to 0MQ / websocket subscribers. + + The RunEngine calls instances of this class with a set of Status objects each time + it enters a wait, and with ``None`` when the wait completes. For each status object + that supports ``watch()``, a callback is registered that streams position/progress + updates. + + Parameters + ---------- + msg_queue : multiprocessing.Queue + Reference to the shared message queue used for publishing messages. + min_update_period : float + Minimum interval in seconds between published updates for a single status + object. The final update (when the status is done) is always sent regardless + of throttling. Default: ``0.2``. + """ + + def __init__(self, *, msg_queue, min_update_period=0.2): + self._msg_queue = msg_queue + self._min_update_period = min_update_period + # Track status objects we have already subscribed to, keyed by id(status) + self._watched = set() + self._last_sent = {} # id(status) -> timestamp of last sent update + self._status_counter = 0 # Counter for generating labels when name is None + + def __call__(self, status_objs_or_none): + """ + Called by the RunEngine with a set of Status objects or ``None``. + """ + if status_objs_or_none is None: + # Waiting is complete — send a completion message and reset state + self._send_completed() + self._watched.clear() + self._last_sent.clear() + self._status_counter = 0 + return + + for st in status_objs_or_none: + st_id = id(st) + if st_id in self._watched: + continue + self._watched.add(st_id) + if not hasattr(st, "watch") or getattr(st, "done", False): + continue + try: + self._status_counter += 1 + label = self._status_counter + st.watch(self._make_callback(st, label)) + except Exception: + logger.debug("Status object does not support watch(): %r", st, exc_info=True) + + def _make_callback(self, status_obj, label): + """ + Create a watch callback bound to a specific status object. + """ + st_id = id(status_obj) + + def _cb( + *, + name=None, + current=None, + initial=None, + target=None, + unit=None, + precision=None, + fraction=None, + time_elapsed=None, + time_remaining=None, + **kwargs, + ): + now = ttime.time() + done = getattr(status_obj, "done", False) + + # The final update (the status reached its target) must never be throttled, + # otherwise progress bars may freeze just short of 100%. ophyd status objects + # emit their last watcher update with ``current == target`` *before* ``done`` + # is set (watchers are cleared once the status settles), and ophyd reports + # ``fraction`` as the fraction *remaining* (0 when the target is reached). + at_target = (fraction is not None and fraction <= 0) or ( + current is not None and target is not None and current == target + ) + is_final = bool(done) or at_target + + # Throttle: skip non-final updates that arrive too quickly + last = self._last_sent.get(st_id, 0) + if not is_final and (now - last) < self._min_update_period: + return + self._last_sent[st_id] = now + + if is_final: + # Clean up tracking for this status + self._last_sent.pop(st_id, None) + + display_name = name if name is not None else f"Status {label}" + + payload = { + "name": display_name, + "current": _to_json_safe(current), + "initial": _to_json_safe(initial), + "target": _to_json_safe(target), + "unit": unit, + "precision": precision, + "fraction": fraction, + "time_elapsed": time_elapsed, + "time_remaining": time_remaining, + "done": is_final, + } + + try: + push_progress_to_msg_queue(msg=payload, msg_queue=self._msg_queue) + except Exception: + logger.debug("Failed to push progress update to msg_queue", exc_info=True) + + return _cb + + def _send_completed(self): + """ + Send a message indicating that the waiting period is complete (all statuses done). + """ + payload = {"completed": True} + try: + push_progress_to_msg_queue(msg=payload, msg_queue=self._msg_queue) + except Exception: + logger.debug("Failed to push progress completion to msg_queue", exc_info=True) diff --git a/src/bluesky_queueserver/manager/start_manager.py b/src/bluesky_queueserver/manager/start_manager.py index fe4d3902..ba4b73b4 100644 --- a/src/bluesky_queueserver/manager/start_manager.py +++ b/src/bluesky_queueserver/manager/start_manager.py @@ -607,9 +607,8 @@ def formatter(prog): dest="zmq_info_addr", type=str, default=None, - help="The address of ZMQ server socket used for publishing information on the state of RE Manager " - "and currently running processes. Currently only the captured STDOUT and STDERR published " - "in 'QS_Console' topic. The parameter overrides the address defined by the environment variable " + help="The address of ZMQ PUB socket used for publishing console output, status info, and " + "progress updates. The parameter overrides the address defined by the environment variable " "'QSERVER_ZMQ_INFO_ADDRESS_FOR_SERVER'. The default address is used if the parameter or the environment " " variable is not defined. Address format: 'tcp://*:60625' " f"(default: {default_zmq_info_address_for_server}).", @@ -624,6 +623,25 @@ def formatter(prog): help="Enable (ON) or disable (OFF) publishing of console output to 0MQ (default: %(default)s).", ) + group_console_output.add_argument( + "--zmq-publish-info", + dest="zmq_publish_info", + type=str, + choices=["ON", "OFF"], + default="ON", + help="Enable (ON) or disable (OFF) publishing of info/status updates to 0MQ (default: %(default)s).", + ) + + group_console_output.add_argument( + "--zmq-publish-progress", + dest="zmq_publish_progress", + type=str, + choices=["ON", "OFF"], + default="ON", + help="Enable (ON) or disable (OFF) publishing of RunEngine waiting/progress updates " + "to 0MQ (default: %(default)s).", + ) + group_console_output.add_argument( "--console-output", dest="console_output", @@ -678,12 +696,15 @@ def formatter(prog): stream_publisher = PublishZMQStreamOutput( msg_queue=msg_queue, console_output_on=settings.print_console_output, - zmq_publish_on=settings.zmq_publish_console, + zmq_publish_console=settings.zmq_publish_console, + zmq_publish_info=settings.zmq_publish_info, + zmq_publish_progress=settings.zmq_publish_progress, zmq_publish_addr=settings.zmq_info_addr, encoding=settings.zmq_encoding, ) - if settings.zmq_publish_console: + zmq_publish_on = settings.zmq_publish_console or settings.zmq_publish_info or settings.zmq_publish_progress + if zmq_publish_on: # Wait for a short period to allow monitoring applications to connect. ttime.sleep(1) @@ -779,6 +800,7 @@ def formatter(prog): config_worker["ipython_control_port"] = settings.ipython_control_port config_worker["ignore_invalid_plans"] = settings.ignore_invalid_plans config_worker["permitted_re_metadata_keys"] = settings.permitted_re_metadata_keys + config_worker["zmq_publish_progress"] = settings.zmq_publish_progress existing_pd_path = settings.existing_plans_and_devices_path if not existing_pd_path: diff --git a/src/bluesky_queueserver/manager/tests/test_info_streaming.py b/src/bluesky_queueserver/manager/tests/test_info_streaming.py index 71f1a70b..9defb7e1 100644 --- a/src/bluesky_queueserver/manager/tests/test_info_streaming.py +++ b/src/bluesky_queueserver/manager/tests/test_info_streaming.py @@ -62,7 +62,7 @@ def test_zmq_info_streaming_1(monkeypatch, re_manager_cmd, stream_enabled): # n params_server = [f"--zmq-info-addr={address_info_server}"] if stream_enabled is not None: - params_server.append(f"--zmq-publish-console={'ON' if stream_enabled else 'OFF'}") + params_server.append(f"--zmq-publish-info={'ON' if stream_enabled else 'OFF'}") zmq_encoding = use_zmq_encoding_for_tests() diff --git a/src/bluesky_queueserver/manager/tests/test_output_streaming.py b/src/bluesky_queueserver/manager/tests/test_output_streaming.py index a33bad40..4dc32840 100644 --- a/src/bluesky_queueserver/manager/tests/test_output_streaming.py +++ b/src/bluesky_queueserver/manager/tests/test_output_streaming.py @@ -118,7 +118,8 @@ def test_ReceiveConsoleOutput_1( pco = PublishZMQStreamOutput( msg_queue=queue, console_output_on=console_output_on, - zmq_publish_on=zmq_publish_on, + zmq_publish_console=zmq_publish_on, + zmq_publish_info=zmq_publish_on, zmq_publish_addr=zmq_publish_addr, zmq_topic_console=zmq_topic_console, zmq_topic_info=zmq_topic_info, @@ -259,7 +260,8 @@ def test_ReceiveConsoleOutputAsync_1(period, cb_type, zmq_encoding, channel): pco = PublishZMQStreamOutput( msg_queue=queue, console_output_on=True, - zmq_publish_on=True, + zmq_publish_console=True, + zmq_publish_info=True, zmq_publish_addr=zmq_publish_addr, zmq_topic_console=zmq_topic_console, zmq_topic_info=zmq_topic_info, @@ -418,7 +420,8 @@ def test_push_info_to_msg_queue_1(zmq_encoding): pco = PublishZMQStreamOutput( msg_queue=queue, console_output_on=True, - zmq_publish_on=True, + zmq_publish_console=True, + zmq_publish_info=True, zmq_publish_addr=zmq_publish_addr, zmq_topic_console=zmq_topic_console, zmq_topic_info=zmq_topic_info, diff --git a/src/bluesky_queueserver/manager/tests/test_plan_monitoring.py b/src/bluesky_queueserver/manager/tests/test_plan_monitoring.py index 3ec08a7a..951601e0 100644 --- a/src/bluesky_queueserver/manager/tests/test_plan_monitoring.py +++ b/src/bluesky_queueserver/manager/tests/test_plan_monitoring.py @@ -2,7 +2,12 @@ import pytest -from bluesky_queueserver.manager.plan_monitoring import CallbackRegisterRun, RunList +from bluesky_queueserver.manager.plan_monitoring import CallbackRegisterRun, RunList, WatcherStreamManager +from bluesky_queueserver.manager.profile_ops import ( + get_default_startup_dir, + load_profile_collection, + load_script_into_existing_nspace, +) def test_RunList_1(): @@ -135,3 +140,306 @@ def test_CallbackRegisterRun_1(scan_id): assert run_list.get_run_list() == [{"uid": uid, "is_open": True, "exit_status": None, **param_expected}] cb("stop", {"run_start": uid, "exit_status": "success"}) assert run_list.get_run_list() == [{"uid": uid, "is_open": False, "exit_status": "success", **param_expected}] + + + +class _MockQueue: + """A simple list-backed mock for multiprocessing.Queue.""" + + def __init__(self): + self.messages = [] + + def put(self, msg): + self.messages.append(msg) + + +class _MockStatus: + """A mock Status object that supports watch().""" + + def __init__(self, *, name="motor1", done=False, supports_watch=True): + self._name = name + self.done = done + self._supports_watch = supports_watch + self._watchers = [] + + def watch(self, func): + if not self._supports_watch: + raise AttributeError("watch not supported") + self._watchers.append(func) + + def simulate_update(self, **kwargs): + """Simulate a watcher callback from the status object.""" + defaults = dict( + name=self._name, + current=None, + initial=None, + target=None, + unit=None, + precision=None, + fraction=None, + time_elapsed=None, + time_remaining=None, + ) + defaults.update(kwargs) + for w in self._watchers: + w(**defaults) + + +def _get_progress_messages(mock_queue): + """Extract only progress payloads from the mock queue.""" + results = [] + for msg in mock_queue.messages: + if msg.get("channel") == "progress" and isinstance(msg.get("msg"), dict): + results.append(msg["msg"]) + return results + + +def test_WatcherStreamManager_basic(): + """ + WatcherStreamManager subscribes to status objects and publishes progress updates. + """ + mq = _MockQueue() + wsm = WatcherStreamManager(msg_queue=mq, min_update_period=0) + + st = _MockStatus(name="motor1") + wsm({st}) + + # Simulate an update + st.simulate_update(current=1.0, initial=0.0, target=5.0, unit="mm", precision=3, fraction=0.2) + + msgs = _get_progress_messages(mq) + assert len(msgs) == 1 + assert msgs[0]["name"] == "motor1" + assert msgs[0]["current"] == 1.0 + assert msgs[0]["target"] == 5.0 + assert msgs[0]["unit"] == "mm" + assert msgs[0]["done"] is False + + # Simulate completion + st.done = True + st.simulate_update(current=5.0, initial=0.0, target=5.0, fraction=1.0) + + msgs = _get_progress_messages(mq) + assert len(msgs) == 2 + assert msgs[1]["done"] is True + assert msgs[1]["current"] == 5.0 + + +def test_WatcherStreamManager_none_clears(): + """ + Calling with None sends a completed message and resets state. + """ + mq = _MockQueue() + wsm = WatcherStreamManager(msg_queue=mq, min_update_period=0) + + st = _MockStatus(name="motor1") + wsm({st}) + st.simulate_update(current=1.0) + + # Signal end of waiting + wsm(None) + + msgs = _get_progress_messages(mq) + # Last message should be the completion indicator + assert msgs[-1] == {"completed": True} + + +def test_WatcherStreamManager_no_resubscribe(): + """ + Calling with the same status object multiple times does not re-subscribe. + """ + mq = _MockQueue() + wsm = WatcherStreamManager(msg_queue=mq, min_update_period=0) + + st = _MockStatus(name="motor1") + wsm({st}) + wsm({st}) + wsm({st}) + + # Should have exactly one watcher registered + assert len(st._watchers) == 1 + + +def test_WatcherStreamManager_throttling(): + """ + Updates that arrive faster than min_update_period are suppressed, + except for the final (done) update. + """ + mq = _MockQueue() + # Set a very long throttle period so all non-final updates are suppressed + wsm = WatcherStreamManager(msg_queue=mq, min_update_period=9999) + + st = _MockStatus(name="motor1") + wsm({st}) + + # First update always goes through (last_sent starts at 0) + st.simulate_update(current=1.0) + # Second update should be throttled + st.simulate_update(current=2.0) + # Third update should be throttled + st.simulate_update(current=3.0) + + msgs = _get_progress_messages(mq) + assert len(msgs) == 1 # Only the first one + + # Final update (done=True) must always go through + st.done = True + st.simulate_update(current=5.0) + + msgs = _get_progress_messages(mq) + assert len(msgs) == 2 + assert msgs[1]["done"] is True + + +def test_WatcherStreamManager_no_watch_support(): + """ + Status objects without watch() are silently ignored. + """ + mq = _MockQueue() + wsm = WatcherStreamManager(msg_queue=mq, min_update_period=0) + + st = _MockStatus(supports_watch=False) + # Should not raise + wsm({st}) + + msgs = _get_progress_messages(mq) + assert len(msgs) == 0 + + +def test_WatcherStreamManager_already_done(): + """ + Status objects that are already done are not subscribed to. + """ + mq = _MockQueue() + wsm = WatcherStreamManager(msg_queue=mq, min_update_period=0) + + st = _MockStatus(done=True) + wsm({st}) + + assert len(st._watchers) == 0 + + +def test_WatcherStreamManager_no_watch_attr(): + """ + Objects without a watch attribute at all are ignored. + """ + mq = _MockQueue() + wsm = WatcherStreamManager(msg_queue=mq, min_update_period=0) + + class _BareStatus: + done = False + + wsm({_BareStatus()}) + msgs = _get_progress_messages(mq) + assert len(msgs) == 0 + + +def test_WatcherStreamManager_multiple_statuses(): + """ + Multiple concurrent status objects each get their own progress stream. + """ + mq = _MockQueue() + wsm = WatcherStreamManager(msg_queue=mq, min_update_period=0) + + st1 = _MockStatus(name="motor1") + st2 = _MockStatus(name="motor2") + wsm({st1, st2}) + + st1.simulate_update(current=1.0, target=5.0) + st2.simulate_update(current=10.0, target=20.0) + + msgs = _get_progress_messages(mq) + assert len(msgs) == 2 + names = {m["name"] for m in msgs} + assert names == {"motor1", "motor2"} + + +def test_WatcherStreamManager_name_none(): + """ + If the watch callback receives name=None, a generated label is used. + """ + mq = _MockQueue() + wsm = WatcherStreamManager(msg_queue=mq, min_update_period=0) + + st = _MockStatus(name=None) + wsm({st}) + st.simulate_update(name=None, current=1.0) + + msgs = _get_progress_messages(mq) + assert len(msgs) == 1 + assert msgs[0]["name"].startswith("Status ") + + +def test_WatcherStreamManager_json_safe_coercion(): + """ + Non-JSON-safe values for current/initial/target are coerced. + """ + import numpy as np + + mq = _MockQueue() + wsm = WatcherStreamManager(msg_queue=mq, min_update_period=0) + + st = _MockStatus(name="motor1") + wsm({st}) + st.simulate_update(current=np.float64(3.14), initial=np.int32(0), target=np.float64(10.0)) + + msgs = _get_progress_messages(mq) + assert len(msgs) == 1 + assert isinstance(msgs[0]["current"], float) + assert isinstance(msgs[0]["initial"], float) + assert isinstance(msgs[0]["target"], float) + + +def test_WatcherStreamManager_sim_motor_move(): + """ + Integration test: load the simulated profile collection, append a script + that creates a slow motor, attach WatcherStreamManager to the RunEngine, + move the motor, and verify that progress updates were published to the queue. + """ + + # Load the simulated profile collection + startup_dir = get_default_startup_dir() + nspace = load_profile_collection(startup_dir, patch_profiles=True) + + # Append a script that creates a slow motor (non-zero delay so watch fires) + script = """ +from ophyd.sim import SynAxis +slow_motor = SynAxis(name="slow_motor", labels={"motors"}) +slow_motor.delay = 0.3 +""" + load_script_into_existing_nspace( + script=script, + nspace=nspace, + script_root_path=startup_dir, + ) + + RE = nspace["RE"] + slow_motor = nspace["slow_motor"] + + # Attach WatcherStreamManager + mq = _MockQueue() + wsm = WatcherStreamManager(msg_queue=mq, min_update_period=0) + RE.waiting_hook = wsm + + # Move the motor from 0 to 1 — this triggers waiting_hook with a MoveStatus + RE(nspace["mv"](slow_motor, 1)) + + msgs = _get_progress_messages(mq) + + # We should have received at least one progress update and one completed message + progress_updates = [m for m in msgs if "completed" not in m] + completed_msgs = [m for m in msgs if m.get("completed") is True] + + assert len(progress_updates) > 0, "Expected at least one progress update during motor move" + assert len(completed_msgs) >= 1, "Expected a 'completed' message after motor move" + + # Verify the structure of a progress update + update = progress_updates[0] + assert "name" in update + assert "current" in update + assert "target" in update + assert update["done"] in (True, False) + assert update["name"] == "slow_motor" + + # The target should be 1 (where we moved the motor) + assert update["target"] == 1 diff --git a/src/bluesky_queueserver/manager/tests/test_start_re_manager_cli.py b/src/bluesky_queueserver/manager/tests/test_start_re_manager_cli.py index e9c3c9fe..8a269ac9 100644 --- a/src/bluesky_queueserver/manager/tests/test_start_re_manager_cli.py +++ b/src/bluesky_queueserver/manager/tests/test_start_re_manager_cli.py @@ -513,6 +513,8 @@ def _get_expected_settings_default_1(_1, _2): "zmq_info_addr": "tcp://*:60625", "zmq_private_key": None, "zmq_publish_console": False, + "zmq_publish_info": True, + "zmq_publish_progress": True, } @@ -616,6 +618,8 @@ def _get_expected_settings_config_2(file_dir, ip_con_dir): "zmq_info_addr": "tcp://*:60627", "zmq_private_key": "Ue=.po0aQ9.}