diff --git a/bec_widgets/utils/qt_data_subscription.py b/bec_widgets/utils/qt_data_subscription.py new file mode 100644 index 00000000..3cb85d88 --- /dev/null +++ b/bec_widgets/utils/qt_data_subscription.py @@ -0,0 +1,114 @@ +""" +Qt bridge for DataAPI subscriptions. + +Wraps a :class:`bec_lib.data_api.Subscription` in a ``QObject``: columnar +:class:`~bec_lib.data_api.SubscriptionUpdate` snapshots arriving on +dispatcher/worker threads are marshalled onto the Qt thread via a queued +signal, stale-scan payloads are dropped, and the subscription is closed with +the widget. This is the one-line integration point for plotting widgets — +no per-widget signal bridges, rate-limit proxies or health checks needed. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from bec_lib.data_api import DataAPI, SourceKey, SubscriptionUpdate +from qtpy.QtCore import QObject, Qt, Signal + +if TYPE_CHECKING: # pragma: no cover + from bec_lib.client import BECClient + + +class QtDataSubscription(QObject): + """One DataAPI subscription delivered on the Qt thread.""" + + #: Emitted on the Qt thread with each SubscriptionUpdate. + updated = Signal(object) + + _raw = Signal(object) + + def __init__( + self, + client: BECClient, + sources: list[SourceKey], + scan: str = "live", + parent: QObject | None = None, + min_emit_interval: float = 0.1, + ): + """ + Subscribe to data for the given sources. + + Args: + client (BECClient): The widget's BEC client. + sources (list[SourceKey]): (device, entry) pairs forming one + correlation group. + scan (str): ``"live"`` to follow the active scan, or a concrete + (possibly finished) scan id. + parent (QObject | None): Qt parent; closing follows the parent's + destruction. + min_emit_interval (float): Backend emission coalescing interval. + + Raises: + ValueError: If a concrete scan id cannot be served. + CorrelationGroupError: If the sources do not form one group. + """ + super().__init__(parent) + self._closed = False + self._subscription = None + # Explicitly queued: the backend delivers the initial backfill + # SYNCHRONOUSLY inside subscribe() when called on the Qt thread; an + # auto-connection would then invoke _filter before _subscription is + # assigned and before the widget had a chance to connect `updated`. + self._raw.connect(self._filter, Qt.QueuedConnection) + self._api = DataAPI(client) + self._subscription = self._api.subscribe( + sources=sources, scan=scan, callback=self._deliver, min_emit_interval=min_emit_interval + ) + self.destroyed.connect(lambda: self.close()) + + # --- api-thread side ----------------------------------------------------- + + def _deliver(self, update: SubscriptionUpdate) -> None: + if not self._closed: + self._raw.emit(update) + + # --- qt-thread side ------------------------------------------------------ + + def _filter(self, update: SubscriptionUpdate) -> None: + if self._closed or self._subscription is None: + return + current = self._subscription.scan_id + if current is not None and update.scan_id != current: + # A payload queued before a rebind; the bound scan's own emission + # follows. + return + self.updated.emit(update) + + # --- public -------------------------------------------------------------- + + @property + def scan_id(self) -> str | None: + """The currently bound scan id.""" + return self._subscription.scan_id + + @property + def sources(self) -> list[SourceKey]: + """The declared source set.""" + return self._subscription.sources + + @property + def healthy(self) -> bool: + """Whether every declared source is currently delivering.""" + return not self._subscription.unbound_sources + + def set_sources(self, sources: list[SourceKey]) -> None: + """Atomically replace the source set.""" + self._subscription.set_sources(sources) + + def close(self) -> None: + """Close the underlying subscription (idempotent).""" + if self._closed: + return + self._closed = True + self._subscription.close() diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index e53063ab..264efd60 100644 --- a/tests/unit_tests/conftest.py +++ b/tests/unit_tests/conftest.py @@ -420,3 +420,46 @@ def _register_worker_thread_dummies(): barrier.wait() except threading.BrokenBarrierError: pass + + +@pytest.fixture(autouse=True) +def _isolate_data_api_instances(): + """ + Drop DataAPI per-client instances between tests. + + The registry is keyed by id(client); mock clients die between tests and + CPython reuses addresses, so a GC-lagged instance from a previous test can + answer for a fresh mock (order-dependent discovery failures). + """ + from bec_lib.data_api import DataAPI + + DataAPI.clear_instance() + yield + DataAPI.clear_instance() + + +@pytest.fixture(autouse=True) +def _isolate_device_signal_info(bec_dispatcher): + """ + Restore the mock devices' ``_info["signals"]`` after each test. + + The dispatcher's fake devices are shared; tests that inject signal + configs (image, multi-waveform, device-input suites) would otherwise + leak them into later tests, making signal discovery order-dependent. + """ + import copy + + devices = getattr(getattr(bec_dispatcher.client, "device_manager", None), "devices", None) + snapshots = {} + if devices is not None: + for name in list(getattr(devices, "keys", dict)() or []): + device = devices[name] + info = getattr(device, "_info", None) + if isinstance(info, dict) and isinstance(info.get("signals"), dict): + snapshots[name] = copy.deepcopy(info["signals"]) + yield + for name, signals in snapshots.items(): + try: + devices[name]._info["signals"] = signals + except (KeyError, TypeError, AttributeError): + continue diff --git a/tests/unit_tests/test_qt_data_subscription.py b/tests/unit_tests/test_qt_data_subscription.py new file mode 100644 index 00000000..94171391 --- /dev/null +++ b/tests/unit_tests/test_qt_data_subscription.py @@ -0,0 +1,100 @@ +"""Tests for the QtDataSubscription bridge.""" + +from unittest import mock + +import pytest +from bec_lib.data_api.models import SubscriptionUpdate + +from bec_widgets.utils.qt_data_subscription import QtDataSubscription + +# pylint: disable=protected-access +# pylint: disable=missing-function-docstring + + +@pytest.fixture +def fake_api(monkeypatch): + subscription = mock.MagicMock() + subscription.scan_id = "scan_1" + subscription.unbound_sources = [] + subscription.sources = [("samx", "samx")] + api = mock.MagicMock() + api.subscribe.return_value = subscription + monkeypatch.setattr("bec_widgets.utils.qt_data_subscription.DataAPI", lambda client: api) + return api, subscription + + +def make_update(scan_id="scan_1"): + return SubscriptionUpdate( + scan_id=scan_id, reason="live", sources={}, aligned_ordinals=(), complete=True + ) + + +def test_updates_are_marshalled_to_qt_thread(qtbot, fake_api): + api, subscription = fake_api + bridge = QtDataSubscription(mock.MagicMock(), sources=[("samx", "samx")]) + received = [] + bridge.updated.connect(received.append) + + callback = api.subscribe.call_args.kwargs["callback"] + update = make_update() + callback(update) + qtbot.waitUntil(lambda: bool(received), timeout=2000) + assert received[0] is update + bridge.close() + assert subscription.close.called + + +def test_stale_scan_payload_dropped(qtbot, fake_api): + api, subscription = fake_api + bridge = QtDataSubscription(mock.MagicMock(), sources=[("samx", "samx")]) + received = [] + bridge.updated.connect(received.append) + callback = api.subscribe.call_args.kwargs["callback"] + + subscription.scan_id = "scan_2" + callback(make_update("scan_1")) # stale + callback(make_update("scan_2")) # current + qtbot.waitUntil(lambda: bool(received), timeout=2000) + assert [u.scan_id for u in received] == ["scan_2"] + bridge.close() + + +def test_health_and_source_delegation(fake_api): + api, subscription = fake_api + bridge = QtDataSubscription(mock.MagicMock(), sources=[("samx", "samx")]) + assert bridge.healthy + subscription.unbound_sources = [("samx", "samx")] + assert not bridge.healthy + bridge.set_sources([("samy", "samy")]) + subscription.set_sources.assert_called_once_with([("samy", "samy")]) + assert bridge.scan_id == "scan_1" + bridge.close() + bridge.close() # idempotent + assert subscription.close.call_count == 1 + + +def test_synchronous_initial_delivery_is_queued(qtbot, monkeypatch): + """The backend delivers the initial backfill synchronously inside + subscribe() on the Qt thread; the bridge must neither crash on its + not-yet-assigned subscription nor lose that first snapshot.""" + subscription = mock.MagicMock() + subscription.scan_id = "scan_1" + + api = mock.MagicMock() + + def synchronous_subscribe(sources, scan, callback, **kwargs): + callback(make_update("scan_1")) + return subscription + + api.subscribe.side_effect = synchronous_subscribe + monkeypatch.setattr("bec_widgets.utils.qt_data_subscription.DataAPI", lambda client: api) + + bridge = QtDataSubscription(mock.MagicMock(), sources=[("samx", "samx")]) + received = [] + # Widgets connect AFTER the constructor returns — the queued initial + # emission must still reach them. + bridge.updated.connect(received.append) + + qtbot.waitUntil(lambda: bool(received), timeout=2000) + assert received[0].scan_id == "scan_1" + bridge.close()