refactor(heatmap): consume the data api for live and history via QtDataSubscription

bec_widgets/widgets/plots/heatmap/heatmap.py

@@ -2,17 +2,16 @@ from __future__ import annotations
 
 import json
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Literal
 
 import numpy as np
 import pyqtgraph as pg
 from bec_lib import bec_logger, messages
-from bec_lib.data_api import DataAPI, DataSubscription
 from bec_lib.endpoints import MessageEndpoints
-from bec_lib.utils.import_utils import lazy_import, lazy_import_from
+from bec_lib.utils.import_utils import lazy_import_from
 from bec_qthemes import material_icon
 from pydantic import BaseModel, Field, field_validator
-from qtpy.QtCore import QObject, QRectF, Qt, QThread, QTimer, Signal
+from qtpy.QtCore import QObject, QRectF, Qt, QThread, Signal
 from qtpy.QtGui import QTransform
 from qtpy.QtWidgets import QDialog, QPushButton, QVBoxLayout
 from toolz import partition
@@ -20,6 +19,7 @@ from toolz import partition
 from bec_widgets.utils.bec_connector import ConnectionConfig
 from bec_widgets.utils.colors import Colors, get_accent_colors
 from bec_widgets.utils.error_popups import SafeProperty, SafeSlot
+from bec_widgets.utils.qt_data_subscription import QtDataSubscription
 from bec_widgets.utils.settings_dialog import SettingsDialog
 from bec_widgets.utils.toolbars.actions import MaterialIconAction
 from bec_widgets.widgets.plots.heatmap.settings.heatmap_setting import HeatmapSettings
@@ -264,10 +264,8 @@ class Heatmap(ImageBase):
 
     new_scan = Signal()
     new_scan_id = Signal(str)
-    sync_signal_update = Signal()
     heatmap_property_changed = Signal()
     interpolation_requested = Signal(object, int)
-    data_api_update = Signal(object, object)
 
     def __init__(self, parent=None, config: HeatmapConfig | None = None, **kwargs):
         if config is None:
@@ -296,8 +294,9 @@ class Heatmap(ImageBase):
         self._interpolation_thread: QThread | None = None
         self._interpolation_worker: _StepInterpolationWorker | None = None
         self._pending_interpolation_request: _InterpolationRequest | None = None
-        self._data_api: DataAPI | None = None
-        self._data_subscription: DataSubscription | None = None
+        self._data_bridge: QtDataSubscription | None = None
+        self._data_bridge_scope: str | None = "live"
+        self._last_columns = None
         self.heatmap_dialog = None
         self.scan_history_dialog = None
         self.scan_history_widget = None
@@ -312,13 +311,16 @@ class Heatmap(ImageBase):
         self.config_label.setVisible(False)
         self.reload = False
         self.bec_dispatcher.connect_slot(self.on_scan_status, MessageEndpoints.scan_status())
-        self.bec_dispatcher.connect_slot(self.on_scan_progress, MessageEndpoints.scan_progress())
-        self.heatmap_property_changed.connect(lambda: self.sync_signal_update.emit())
-        self.data_api_update.connect(self.update_plot)
-
-        self.proxy_update_sync = pg.SignalProxy(
-            self.sync_signal_update, rateLimit=5, slot=self.update_plot
-        )
+        # Display-property changes re-render the cached columns; all data
+        # flows through the DataAPI subscription. Queued: the interpolation
+        # worker expects requests from event-loop ticks, and back-to-back
+        # property changes coalesce per tick instead of re-entering
+        # update_plot inside a property setter.
+        self.heatmap_property_changed.connect(self.update_plot, Qt.QueuedConnection)
+        if self._config_sources() is not None:
+            # Restored from a saved configuration: start the data feed without
+            # requiring a plot() call.
+            self._start_data_feed()
         self._init_toolbar_heatmap()
         self.toolbar.show_bundles(
             [
@@ -463,11 +465,22 @@ class Heatmap(ImageBase):
             return
 
         self._history_scan_id = None
-        self._fetch_running_scan()
-        self._setup_data_api_subscription()
-        # Also notifies settings widgets and triggers a plot update via sync_signal_update
+        self._start_data_feed()
+        # Also notifies settings widgets, which re-render via update_plot
         self.heatmap_property_changed.emit()
 
+    def _start_data_feed(self):
+        """
+        Bind the widget to the running scan or, while idle, to the latest
+        finished scan (served by the history plugin); a live-follow bridge
+        replaces the idle binding when the next scan starts.
+        """
+        self._fetch_running_scan()
+        if self.scan_item is not None and not hasattr(self.scan_item, "live_data"):
+            self._setup_data_api_subscription(scan=self.scan_id)
+        else:
+            self._setup_data_api_subscription()
+
     def _fetch_running_scan(self):
         scan = self.client.queue.scan_storage.current_scan
         if scan is not None:
@@ -526,7 +539,14 @@ class Heatmap(ImageBase):
         # plot() is called again without a scan_id.
         self._history_scan_id = self.scan_id
 
-        # Also notifies settings widgets and triggers a plot update via sync_signal_update
+        # Fetch the history data through the DataAPI (file-backed, worker
+        # thread).
+        try:
+            self._setup_data_api_subscription(scan=self.scan_id)
+        except Exception as exc:  # pylint: disable=broad-except
+            logger.warning(f"History data-api subscription failed: {exc}")
+
+        # Also notifies settings widgets, which re-render via update_plot
         self.heatmap_property_changed.emit()
 
     def update_labels(self):
@@ -713,90 +733,45 @@ class Heatmap(ImageBase):
         self.toolbar.components.get_action("heatmap_settings").action.setChecked(False)
 
     def _cleanup_data_api_subscription(self):
-        if self._data_subscription is None:
+        if self._data_bridge is None:
             return
         try:
-            self._data_subscription.close()
+            self._data_bridge.close()
         finally:
-            self._data_subscription = None
+            self._data_bridge = None
 
-    def _setup_data_api_subscription(self):
-        self._cleanup_data_api_subscription()
-
-        if self._history_scan_id is not None or self._image_config is None:
-            return
-
-        if not all(
-            [
-                self._image_config.device_x,
-                self._image_config.device_y,
-                self._image_config.device_z,
+    def _config_sources(self):
+        if self._image_config is None:
+            return None
+        try:
+            return [
+                (self._image_config.device_x.device, self._image_config.device_x.signal),
+                (self._image_config.device_y.device, self._image_config.device_y.signal),
+                (self._image_config.device_z.device, self._image_config.device_z.signal),
             ]
-        ):
+        except AttributeError:
+            return None
+
+    def _setup_data_api_subscription(self, scan: str = "live"):
+        self._cleanup_data_api_subscription()
+        self._data_bridge_scope = scan
+
+        if scan == "live" and self._history_scan_id is not None:
+            return
+        sources = self._config_sources()
+        if sources is None or not all(dev for dev, _ in sources):
             return
 
         try:
-            if self._data_api is None:
-                self._data_api = DataAPI(self.client)
-            subscription = self._data_api.create_subscription(live=True, buffered=True)
-            subscription.set_callback(self.data_api_update.emit)
-            subscription.add_device(
-                self._image_config.device_x.device, self._image_config.device_x.signal
-            )
-            subscription.add_device(
-                self._image_config.device_y.device, self._image_config.device_y.signal
-            )
-            subscription.add_device(
-                self._image_config.device_z.device, self._image_config.device_z.signal
+            self._data_bridge = QtDataSubscription(
+                self.client, sources=sources, scan=scan, parent=self, min_emit_interval=0.2
             )
+            self._data_bridge.updated.connect(self._on_data_update)
         except Exception as exc:
             logger.warning(f"Failed to configure heatmap data-api subscription: {exc}")
             self._cleanup_data_api_subscription()
             return
 
-        self._data_subscription = subscription
-
-    @staticmethod
-    def _normalize_series_values(values):
-        if values is None:
-            return None
-        if isinstance(values, np.ndarray):
-            return values.tolist()
-        if isinstance(values, list):
-            return values
-        if isinstance(values, tuple):
-            return list(values)
-        return [values]
-
-    def _extract_buffered_series(
-        self, data: dict, device: str, signal: str
-    ) -> list[Any] | None:
-        signal_buffer = data.get(device, {}).get(signal)
-        if signal_buffer is None:
-            return None
-        if isinstance(signal_buffer, dict):
-            return self._normalize_series_values(signal_buffer.get("value"))
-        if not isinstance(signal_buffer, list):
-            return None
-
-        values = []
-        for item in signal_buffer:
-            if not isinstance(item, dict) or "value" not in item:
-                return None
-            values.append(item["value"])
-        return values
-
-    def _extract_scan_series(self, data, access_key: str, device: str, signal: str) -> list[Any] | None:
-        if access_key == "val":
-            values = data.get(device, {}).get(signal, {}).get(access_key, None)
-            return self._normalize_series_values(values)
-
-        readback = data.get(device, {}).get(signal, None)
-        if readback is None:
-            return None
-        values = readback.read().get("value", None)
-        return self._normalize_series_values(values)
-
     @SafeSlot(dict, dict)
     def on_scan_status(self, msg: dict, meta: dict):
         """
@@ -819,57 +794,55 @@ class Heatmap(ImageBase):
             self.old_scan_id = self.scan_id
             self.scan_id = current_scan_id
             self.scan_item = self.queue.scan_storage.find_scan_by_ID(self.scan_id)  # type: ignore
+            if self._data_bridge is None or self._data_bridge_scope != "live":
+                # The widget started idle and bound to the latest finished
+                # scan (or has no bridge yet); a scan is running now, so
+                # switch to a live-follow subscription.
+                self._setup_data_api_subscription()
 
-            if self._data_subscription is None:
-                # First trigger to update the scan curves
-                self.sync_signal_update.emit()
-
-    @SafeSlot(dict, dict)
-    def on_scan_progress(self, msg: dict, meta: dict):
-        if self._history_scan_id is not None:
-            return
-        if self._data_subscription is not None:
-            return
-        self.sync_signal_update.emit()
-        status = msg.get("done")
-        if status:
-            QTimer.singleShot(100, self.update_plot)
-            QTimer.singleShot(300, self.update_plot)
-
-    @SafeSlot(verify_sender=True)
-    def update_plot(self, data: dict | None = None, metadata: dict | None = None) -> None:
+    @SafeSlot(object)
+    def _on_data_update(self, update) -> None:
         """
-        Update the plot with the current data.
-        """
-        if self.scan_item is None:
-            logger.info("No scan executed so far; skipping update.")
-            return
+        Render one columnar DataAPI update (live or history).
 
-        if self._image_config is None:
+        Args:
+            update (SubscriptionUpdate): Aligned full-state snapshot.
+        """
+        sources = self._config_sources()
+        if sources is None or self._data_bridge is None:
             return
+        if self._history_scan_id is not None and update.scan_id != self._history_scan_id:
+            return
+        columns = update.aligned()
         try:
-            device_x = self._image_config.device_x.device
-            signal_x = self._image_config.device_x.signal
-            device_y = self._image_config.device_y.device
-            signal_y = self._image_config.device_y.signal
-            device_z = self._image_config.device_z.device
-            signal_z = self._image_config.device_z.signal
-        except AttributeError:
+            x_data = list(columns[tuple(sources[0])])
+            y_data = list(columns[tuple(sources[1])])
+            z_data = list(columns[tuple(sources[2])])
+        except KeyError:
             return
+        if not x_data:
+            return
+        self._render_columns(x_data, y_data, z_data)
 
-        if isinstance(data, dict):
-            x_data = self._extract_buffered_series(data, device_x, signal_x)
-            y_data = self._extract_buffered_series(data, device_y, signal_y)
-            z_data = self._extract_buffered_series(data, device_z, signal_z)
-        else:
-            data, access_key = self._fetch_scan_data_and_access()
-            if data == "none":
-                logger.info("No scan executed so far; skipping update.")
-                return
-            x_data = self._extract_scan_series(data, access_key, device_x, signal_x)
-            y_data = self._extract_scan_series(data, access_key, device_y, signal_y)
-            z_data = self._extract_scan_series(data, access_key, device_z, signal_z)
+    @SafeSlot()
+    def update_plot(self, _=None) -> None:
+        """
+        Re-render the last received data columns.
 
+        All data flows through the DataAPI subscription; this slot only
+        re-applies the most recent columns after a display-property change
+        (color map, interpolation, labels).
+        """
+        if self._last_columns is None:
+            return
+        x_data, y_data, z_data = self._last_columns
+        self._render_columns(x_data, y_data, z_data)
+
+    def _render_columns(self, x_data, y_data, z_data) -> None:
+        """
+        Render one aligned x/y/z column set into the heatmap image.
+        """
+        self._last_columns = (x_data, y_data, z_data)
         if x_data is None or y_data is None or z_data is None:
             logger.warning("x, y, or z data is None; skipping update.")
             return
@@ -883,17 +856,17 @@ class Heatmap(ImageBase):
         if hasattr(self.scan_item, "status_message"):
             scan_msg = self.scan_item.status_message
         elif hasattr(self.scan_item, "metadata"):
-            metadata = self.scan_item.metadata["bec"]
-            status = metadata["status"]
-            scan_id = metadata["scan_id"]
-            scan_name = metadata["scan_name"]
-            scan_type = metadata["scan_type"]
-            scan_number = metadata["scan_number"]
-            request_inputs = metadata["request_inputs"]
+            bec_metadata = self.scan_item.metadata["bec"]
+            status = bec_metadata["status"]
+            scan_id = bec_metadata["scan_id"]
+            scan_name = bec_metadata["scan_name"]
+            scan_type = bec_metadata["scan_type"]
+            scan_number = bec_metadata["scan_number"]
+            request_inputs = bec_metadata["request_inputs"]
             if "arg_bundle" in request_inputs and isinstance(request_inputs["arg_bundle"], str):
                 # Convert the arg_bundle from a JSON string to a dictionary
                 request_inputs["arg_bundle"] = json.loads(request_inputs["arg_bundle"])
-            positions = metadata.get("positions", [])
+            positions = bec_metadata.get("positions", [])
             positions = positions.tolist() if isinstance(positions, np.ndarray) else positions
 
             scan_msg = messages.ScanStatusMessage(
@@ -1394,30 +1367,6 @@ class Heatmap(ImageBase):
             params[cmds[0]] = list(cmds[1:])
         return params
 
-    def _fetch_scan_data_and_access(self):
-        """
-        Decide whether the widget is in live or historical mode
-        and return the appropriate data dict and access key.
-
-        Returns:
-            data_dict (dict): The data structure for the current scan.
-            access_key (str): Either 'val' (live) or 'value' (history).
-        """
-        if self.scan_item is None:
-            # Optionally fetch the latest from history if nothing is set
-            # self.update_with_scan_history(-1)
-            if self.scan_item is None:
-                logger.info("No scan executed so far; skipping update.")
-                return "none", "none"
-
-        if hasattr(self.scan_item, "live_data"):
-            # Live scan
-            return self.scan_item.live_data, "val"
-
-        # Historical
-        scan_devices = self.scan_item.devices
-        return scan_devices, "value"
-
     def reset(self):
         self._cancel_interpolation()
         self._grid_index = None

tests/unit_tests/test_heatmap_widget.py

@@ -30,33 +30,33 @@ def heatmap_widget(qtbot, mocked_client):
     yield widget
 
 
-class _FakeDataSubscription:
-    def __init__(self):
-        self.callback = None
-        self.devices = []
+class _FakeBridge:
+    """Stand-in for QtDataSubscription."""
+
+    def __init__(self, client, sources, scan="live", parent=None, min_emit_interval=0.1):
+        self.client = client
+        self.sources = list(sources)
+        self.scan = scan
+        self.min_emit_interval = min_emit_interval
+        self.healthy = True
         self.closed = False
-
-    def set_callback(self, callback):
-        self.callback = callback
-        return self
-
-    def add_device(self, device, signal):
-        self.devices.append((device, signal))
-        return self
+        self.updated = mock.MagicMock()
 
     def close(self):
         self.closed = True
 
 
-class _FakeDataAPI:
-    def __init__(self, client):
-        self.client = client
-        self.create_subscription_calls = []
-        self.subscription = _FakeDataSubscription()
+@pytest.fixture
+def fake_bridges(monkeypatch):
+    created = []
 
-    def create_subscription(self, **kwargs):
-        self.create_subscription_calls.append(kwargs)
-        return self.subscription
+    def factory(client, sources, scan="live", parent=None, min_emit_interval=0.1):
+        bridge = _FakeBridge(client, sources, scan=scan, min_emit_interval=min_emit_interval)
+        created.append(bridge)
+        return bridge
+
+    monkeypatch.setattr("bec_widgets.widgets.plots.heatmap.heatmap.QtDataSubscription", factory)
+    return created
 
 
 def test_heatmap_plot(heatmap_widget):
@@ -67,12 +67,7 @@ def test_heatmap_plot(heatmap_widget):
     assert heatmap_widget._image_config.device_z.device == "bpm4i"
 
 
-def test_heatmap_plot_sets_up_live_data_api_subscription(heatmap_widget, monkeypatch):
-    fake_data_api = _FakeDataAPI(heatmap_widget.client)
-    monkeypatch.setattr(
-        "bec_widgets.widgets.plots.heatmap.heatmap.DataAPI", lambda client: fake_data_api
-    )
-
+def test_heatmap_plot_sets_up_live_data_api_subscription(heatmap_widget, fake_bridges):
     heatmap_widget.plot(
         device_x="samx",
         device_y="samy",
@@ -82,14 +77,12 @@ def test_heatmap_plot_sets_up_live_data_api_subscription(heatmap_widget, monkeyp
         signal_z="bpm4i",
     )
 
-    assert fake_data_api.create_subscription_calls == [{"live": True, "buffered": True}]
-    assert fake_data_api.subscription.callback == heatmap_widget.data_api_update.emit
-    assert fake_data_api.subscription.devices == [
-        ("samx", "samx"),
-        ("samy", "samy"),
-        ("bpm4i", "bpm4i"),
-    ]
-    assert heatmap_widget._data_subscription is fake_data_api.subscription
+    assert len(fake_bridges) == 1
+    bridge = fake_bridges[0]
+    assert bridge.scan == "live"
+    assert bridge.sources == [("samx", "samx"), ("samy", "samy"), ("bpm4i", "bpm4i")]
+    assert bridge.updated.connect.called
+    assert heatmap_widget._data_bridge is bridge
 
 
 def test_heatmap_plot_with_scan_id_uses_history(heatmap_widget):
@@ -140,15 +133,39 @@ def test_heatmap_update_with_scan_history_resets_cached_image_state(heatmap_widg
     assert heatmap_widget.scan_id == "scan-456"
 
 
-def test_heatmap_update_with_scan_history_closes_live_data_api_subscription(heatmap_widget):
+def test_heatmap_update_with_scan_history_closes_live_data_api_subscription(
+    heatmap_widget, fake_bridges
+):
     history_scan = mock.MagicMock()
     history_scan.scan_id = "scan-456"
-    heatmap_widget._data_subscription = _FakeDataSubscription()
+    live_bridge = _FakeBridge(heatmap_widget.client, [("samx", "samx")])
+    heatmap_widget._data_bridge = live_bridge
 
     with mock.patch.object(heatmap_widget, "get_history_scan_item", return_value=history_scan):
         heatmap_widget.update_with_scan_history(scan_id="scan-456")
 
-    assert heatmap_widget._data_subscription is None
+    assert live_bridge.closed is True
+    # Without an image config no history subscription is created.
+    assert heatmap_widget._data_bridge is None
+
+
+def test_heatmap_update_with_scan_history_subscribes_to_history_scan(heatmap_widget, fake_bridges):
+    history_scan = mock.MagicMock()
+    history_scan.scan_id = "scan-456"
+    heatmap_widget.plot(
+        device_x="samx",
+        device_y="samy",
+        device_z="bpm4i",
+        signal_x="samx",
+        signal_y="samy",
+        signal_z="bpm4i",
+    )
+    with mock.patch.object(heatmap_widget, "get_history_scan_item", return_value=history_scan):
+        heatmap_widget.update_with_scan_history(scan_id="scan-456")
+
+    assert fake_bridges[-1].scan == "scan-456"
+    assert heatmap_widget._data_bridge is fake_bridges[-1]
+    assert heatmap_widget._history_scan_id == "scan-456"
 
 
 def test_heatmap_on_scan_status_resets_after_history_scan_selection(heatmap_widget):
@@ -427,6 +444,8 @@ def test_heatmap_update_plot_no_scan_item(heatmap_widget):
 
 
 def test_heatmap_update_plot(heatmap_widget):
+    """update_plot re-renders the most recent data columns (display-property
+    changes); it performs no data fetching of its own."""
     heatmap_widget._image_config = HeatmapConfig(
         parent_id="parent_id",
         device_x=HeatmapDeviceSignal(device="samx", signal="samx"),
@@ -447,13 +466,23 @@ def test_heatmap_update_plot(heatmap_widget):
         },
         request_inputs={"arg_bundle": ["samx", -5, 5, 10, "samy", -5, 5, 10], "kwargs": {}},
     )
+    heatmap_widget.status_message = heatmap_widget.scan_item.status_message
+
+    # Nothing cached yet: a re-render request is a no-op.
     with mock.patch.object(heatmap_widget.main_image, "setImage") as mock_set_image:
-        heatmap_widget.update_plot(_override_slot_params={"verify_sender": False})
+        heatmap_widget.update_plot()
+    mock_set_image.assert_not_called()
+
+    heatmap_widget._last_columns = (list(x_levels), list(y_levels), [float(i) for i in range(10)])
+    with mock.patch.object(heatmap_widget.main_image, "setImage") as mock_set_image:
+        heatmap_widget.update_plot()
         img = mock_set_image.mock_calls[0].args[0]
         assert img.shape == (10, 10)
 
 
-def test_heatmap_update_plot_from_buffered_data_api_payload(heatmap_widget):
+def test_heatmap_renders_columns_from_data_api_update(heatmap_widget):
+    from bec_lib.data_api.models import SourceData, SubscriptionUpdate
+
     heatmap_widget._image_config = HeatmapConfig(
         parent_id="parent_id",
         device_x=HeatmapDeviceSignal(device="samx", signal="samx"),
@@ -461,6 +490,7 @@ def test_heatmap_update_plot_from_buffered_data_api_payload(heatmap_widget):
         device_z=HeatmapDeviceSignal(device="bpm4i", signal="bpm4i"),
         color_map="viridis",
     )
+    heatmap_widget._data_bridge = _FakeBridge(heatmap_widget.client, [("samx", "samx")])
     heatmap_widget.scan_item = create_dummy_scan_item()
     x_levels = np.linspace(-5, 5, 10).tolist()
     y_levels = np.linspace(-5, 5, 10).tolist()
@@ -474,20 +504,38 @@ def test_heatmap_update_plot_from_buffered_data_api_payload(heatmap_widget):
         },
         request_inputs={"arg_bundle": ["samx", -5, 5, 10, "samy", -5, 5, 10], "kwargs": {}},
     )
-    payload = {
-        "samx": {"samx": [{"value": value, "timestamp": idx} for idx, value in enumerate(x_levels)]},
-        "samy": {"samy": [{"value": value, "timestamp": idx} for idx, value in enumerate(y_levels)]},
-        "bpm4i": {
-            "bpm4i": [{"value": idx, "timestamp": idx} for idx in range(len(x_levels))]
+    heatmap_widget.status_message = heatmap_widget.scan_item.status_message
+
+    n = 10
+    xs = tuple(x_levels)
+    ys = tuple(y_levels)
+    zs = tuple(float(i) for i in range(n))
+
+    def source(dev, values):
+        return SourceData(
+            device=dev,
+            entry=dev,
+            kind="monitored",
+            ordinals=tuple(range(n)),
+            values=values,
+            timestamps=tuple(float(i) for i in range(n)),
+            complete=True,
+        )
+
+    update = SubscriptionUpdate(
+        scan_id="123",
+        reason="live",
+        sources={
+            ("samx", "samx"): source("samx", xs),
+            ("samy", "samy"): source("samy", ys),
+            ("bpm4i", "bpm4i"): source("bpm4i", zs),
         },
-    }
+        aligned_ordinals=tuple(range(n)),
+        complete=True,
+    )
 
     with mock.patch.object(heatmap_widget.main_image, "setImage") as mock_set_image:
-        heatmap_widget.update_plot(
-            data=payload,
-            metadata={"scan_id": "123"},
-            _override_slot_params={"verify_sender": False},
-        )
+        heatmap_widget._on_data_update(update)
         img = mock_set_image.mock_calls[0].args[0]
         assert img.shape == (10, 10)
 
@@ -1057,8 +1105,8 @@ def test_device_properties_with_none_values(heatmap_widget):
 
 def test_heatmap_history_mode_ignores_live_scan_updates(heatmap_widget):
     """
-    Once pinned to a history scan, the heatmap ignores live scan status/progress updates
-    until plot() is called again without a scan_id.
+    Once pinned to a history scan, live scan-status updates neither reset the
+    widget nor steal the pin until plot() is called again without a scan_id.
     """
     history_scan = mock.MagicMock()
     history_scan.scan_id = "scan-456"
@@ -1074,10 +1122,6 @@ def test_heatmap_history_mode_ignores_live_scan_updates(heatmap_widget):
     reset_mock.assert_not_called()
     assert heatmap_widget.scan_id == "scan-456"
 
-    with mock.patch.object(heatmap_widget, "sync_signal_update") as sync_mock:
-        heatmap_widget.on_scan_progress({"done": False}, {})
-    sync_mock.emit.assert_not_called()
-
     # Plotting without a scan_id returns to live mode
     heatmap_widget.plot(device_x="samx", device_y="samy", device_z="bpm4i")
     assert heatmap_widget._history_scan_id is None
@@ -1220,3 +1264,120 @@ def test_heatmap_settings_scan_index_syncs_with_widget(heatmap_widget, qtbot, sc
 
     heatmap_widget.heatmap_dialog.reject()
     qtbot.waitUntil(lambda: heatmap_widget.heatmap_dialog is None)
+
+
+def test_heatmap_has_no_legacy_data_path(heatmap_widget):
+    """All data flows through the DataAPI: the legacy fetch machinery and its
+    scan-progress trigger no longer exist on the widget."""
+    for legacy in (
+        "sync_signal_update",
+        "on_scan_progress",
+        "_fetch_scan_data_and_access",
+        "_extract_scan_series",
+        "_data_api_feed_healthy",
+        "proxy_update_sync",
+    ):
+        assert not hasattr(heatmap_widget, legacy)
+
+
+def test_heatmap_on_data_update_respects_history_pin(heatmap_widget):
+    """While pinned to a history scan, only that scan's updates render."""
+    from bec_lib.data_api.models import SourceData, SubscriptionUpdate
+
+    heatmap_widget.plot(
+        device_x="samx",
+        device_y="samy",
+        device_z="bpm4i",
+        signal_x="samx",
+        signal_y="samy",
+        signal_z="bpm4i",
+    )
+    heatmap_widget._data_bridge = _FakeBridge(heatmap_widget.client, [("samx", "samx")])
+    heatmap_widget._history_scan_id = "scan-hist"
+
+    def update_for(scan_id):
+        sources = {}
+        for dev in ("samx", "samy", "bpm4i"):
+            sources[(dev, dev)] = SourceData(
+                device=dev,
+                entry=dev,
+                kind="monitored",
+                ordinals=(0,),
+                values=(1.0,),
+                timestamps=(1.0,),
+                complete=True,
+            )
+        return SubscriptionUpdate(
+            scan_id=scan_id, reason="history", sources=sources, aligned_ordinals=(0,), complete=True
+        )
+
+    with mock.patch.object(heatmap_widget, "_render_columns") as render:
+        heatmap_widget._on_data_update(update_for("scan-live"))
+    render.assert_not_called()
+
+    with mock.patch.object(heatmap_widget, "_render_columns") as render:
+        heatmap_widget._on_data_update(update_for("scan-hist"))
+    render.assert_called_once_with([1.0], [1.0], [1.0])
+
+
+def test_heatmap_failed_subscription_setup_leaves_no_bridge(heatmap_widget, monkeypatch):
+    """A failing subscription constructor must leave no half-configured feed."""
+
+    def raising_factory(*args, **kwargs):
+        raise ValueError("not bundle compatible")
+
+    monkeypatch.setattr(
+        "bec_widgets.widgets.plots.heatmap.heatmap.QtDataSubscription", raising_factory
+    )
+    heatmap_widget.plot(
+        device_x="samx",
+        device_y="samy",
+        device_z="bpm4i",
+        signal_x="samx",
+        signal_y="samy",
+        signal_z="bpm4i",
+    )
+    assert heatmap_widget._data_bridge is None
+
+
+def test_heatmap_switches_history_bound_bridge_to_live_on_new_scan(heatmap_widget, fake_bridges):
+    """A widget configured while idle binds to the latest finished scan; when
+    a new scan starts it must switch to a live-follow subscription (live
+    regression: the history-bound bridge froze the plot for the whole scan)."""
+    heatmap_widget.plot(
+        device_x="samx",
+        device_y="samy",
+        device_z="bpm4i",
+        signal_x="samx",
+        signal_y="samy",
+        signal_z="bpm4i",
+    )
+    # Simulate the idle-startup state: bridge bound to a finished scan.
+    heatmap_widget._setup_data_api_subscription(scan="old-finished-scan")
+    assert fake_bridges[-1].scan == "old-finished-scan"
+
+    heatmap_widget.on_scan_status({"scan_id": "new-live-scan"}, {})
+
+    assert fake_bridges[-1].scan == "live"
+    assert heatmap_widget._data_bridge is fake_bridges[-1]
+
+    # A further scan status for the live-follow bridge must NOT rebuild it.
+    bridge_count = len(fake_bridges)
+    heatmap_widget.on_scan_status({"scan_id": "another-scan"}, {})
+    assert len(fake_bridges) == bridge_count
+
+
+def test_heatmap_restored_from_config_starts_data_feed(qtbot, mocked_client, fake_bridges):
+    """A heatmap recreated from a saved configuration (no plot() call) must
+    start its DataAPI feed just like a freshly plotted one."""
+    config = HeatmapConfig(
+        parent_id="parent_id",
+        device_x=HeatmapDeviceSignal(device="samx", signal="samx"),
+        device_y=HeatmapDeviceSignal(device="samy", signal="samy"),
+        device_z=HeatmapDeviceSignal(device="bpm4i", signal="bpm4i"),
+        color_map="plasma",
+    )
+    widget = Heatmap(client=mocked_client, config=config)
+    qtbot.addWidget(widget)
+    assert widget._data_bridge is not None
+    assert fake_bridges[-1].sources == [("samx", "samx"), ("samy", "samy"), ("bpm4i", "bpm4i")]