fix(test): update last scan data retrieval in async plotting test

bec_widgets/widgets/plots/waveform/waveform.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import traceback
 from typing import Literal
 
 import lmfit
@@ -1205,6 +1206,7 @@ class Waveform(PlotBase):
         plot_mode = self.x_axis_mode["name"]
         for curve in self._async_curves:
             x_data = None  # Reset x_data
+            y_data = None
             # Get the curve data
             async_data = msg["signals"].get(curve.config.signal.entry, None)
             if async_data is None:
@@ -1223,6 +1225,7 @@ class Waveform(PlotBase):
                         data_plot_y = data_plot_y[-1, :]
                 else:
                     x_data, y_data = curve.get_data()
+
                     if y_data is not None:
                         data_plot_y = np.hstack((y_data, data_plot_y))
             # Add slice

tests/end-2-end/test_plotting_framework_e2e.py

@@ -3,11 +3,14 @@ import time
 import numpy as np
 import pytest
 from bec_lib.endpoints import MessageEndpoints
+from bec_lib.logger import bec_logger
 
 from bec_widgets.cli.client import Image, MotorMap, MultiWaveform, ScatterWaveform, Waveform
 from bec_widgets.cli.rpc.rpc_base import RPCReference
 from bec_widgets.tests.utils import check_remote_data_size
 
+logger = bec_logger.logger
+
 
 def test_rpc_waveform1d_custom_curve(qtbot, connected_client_gui_obj):
     gui = connected_client_gui_obj
@@ -111,6 +114,7 @@ def test_rpc_waveform_scan(qtbot, bec_client_lib, connected_client_gui_obj):
 
 
 def test_async_plotting(qtbot, bec_client_lib, connected_client_gui_obj):
+    logger.warning("Starting async plotting test in e2e test.")
     gui = connected_client_gui_obj
     dock = gui.bec
 
@@ -121,9 +125,9 @@ def test_async_plotting(qtbot, bec_client_lib, connected_client_gui_obj):
 
     # Test add
     dev.waveform.sim.select_model("GaussianModel")
-    dev.waveform.sim.params = {"amplitude": 1000, "center": 4000, "sigma": 300}
+    dev.waveform.sim.params = {"amplitude": 1000, "center": 400, "sigma": 300}
     dev.waveform.async_update.put("add")
-    dev.waveform.waveform_shape.put(10000)
+    dev.waveform.waveform_shape.put(1000)  # Do not reduce, data needs to be large to downsample
     wf = dock.new("wf_dock").new("Waveform")
     curve = wf.plot(y_name="waveform")
 
@@ -131,28 +135,24 @@ def test_async_plotting(qtbot, bec_client_lib, connected_client_gui_obj):
     status.wait()
 
     # Wait for the scan to finish and the data to be available in history
-    # Wait until scan_id is in history
     def _wait_for_scan_in_history():
-        if len(client.history) == 0:
-            return False
-        # Once items appear in storage, the last one hast to be the one we just scanned
-        return client.history[-1].metadata.bec["scan_id"] == status.scan.scan_id
+        # Get scan item from history
+        scan_item = client.history.get_by_scan_id(status.scan.scan_id)
+        return scan_item is not None
+
+    qtbot.waitUntil(_wait_for_scan_in_history, timeout=7000)
+    # Get all data
+    waveform_data = client.history[-1].devices.waveform.waveform_waveform.read()["value"]
 
-    qtbot.waitUntil(_wait_for_scan_in_history, timeout=10000)
-    last_scan_data = client.history[-1]
     # check plotted data
     x_data, y_data = curve.get_data()
     assert np.array_equal(x_data, np.linspace(0, len(y_data) - 1, len(y_data)))
-    assert np.array_equal(
-        y_data, last_scan_data.devices.waveform.get("waveform_waveform", {}).read().get("value", [])
-    )
+    assert np.array_equal(y_data, waveform_data)
 
     # Check displayed data
     x_data_display, y_data_display = curve._get_displayed_data()
     # Should be not more than 1% difference, actually be closer but this might be flaky
     assert np.isclose(x_data_display[-1], x_data[-1], rtol=0.01)
-    # Downsampled data should be smaller than original data
-    assert len(y_data_display) < len(y_data)
 
 
 def test_rpc_image(qtbot, bec_client_lib, connected_client_gui_obj):

tests/end-2-end/test_rpc_widgets_e2e.py

@@ -315,6 +315,14 @@ def test_widgets_e2e_image(
     # Run a scan and plot the image
     s = scans.line_scan(dev.samx, -3, 3, steps=50, exp_time=0.01, relative=False)
     s.wait()
+
+    def _wait_for_scan_in_history():
+        # Get scan item from history
+        scan_item = bec.history.get_by_scan_id(s.scan.scan_id)
+        return scan_item is not None
+
+    qtbot.waitUntil(_wait_for_scan_in_history, timeout=7000)
+
     # Check that last image is equivalent to data in Redis
     last_img = bec.device_monitor.get_data(
         dev.eiger, count=1
@@ -325,22 +333,6 @@ def test_widgets_e2e_image(
     maybe_remove_widget(qtbot, gui, dock, widget, random_generator_from_seed)
     maybe_remove_dock_area(qtbot, gui, dock_area, random_generator_from_seed)
 
-    # @pytest.mark.timeout(PYTEST_TIMEOUT)
-    # def test_widgets_e2e_lmfit_dialog(
-    #     qtbot, connected_gui_and_bec_with_scope_session, random_generator_from_seed
-    # ):
-    #     """Test the LMFITDialog widget."""
-    #     gui = connected_gui_and_bec_with_scope_session
-    #     bec = gui._client
-    #     # Create dock_area, dock, widget
-    #     dock_area, dock, widget = create_widget(qtbot, gui, gui.available_widgets.LMFitDialog)
-
-    #     # No rpc calls to check so far
-
-    #     # Test removing the widget, or leaving it open for the next test
-    #     maybe_remove_widget(qtbot, gui, dock, widget, random_generator_from_seed)
-    maybe_remove_dock_area(qtbot, gui, dock_area, random_generator_from_seed)
-
 
 @pytest.mark.timeout(PYTEST_TIMEOUT)
 def test_widgets_e2e_log_panel(
@@ -672,24 +664,16 @@ def test_widgets_e2e_waveform(
     s = scans.line_scan(dev.samx, -3, 3, steps=50, exp_time=0.01, relative=False)
     s.wait()
 
-    def _wait_live_data_updated():
-        # Wait until storage exists
-        if len(bec.queue.scan_storage.storage) == 0:
-            return False
-        scan_item = bec.queue.scan_storage.storage[-1]
-        # Wait until scan_id is in history
-        if not scan_item.status_message.info["scan_id"] == s.scan.scan_id:
-            return False
-        # Wait until data for all steps is available
-        return len(scan_item.live_data.samx.samx.val) == 50
+    def _wait_for_scan_in_history():
+        # Get scan item from history
+        scan_item = bec.history.get_by_scan_id(s.scan.scan_id)
+        return scan_item is not None
 
-    qtbot.waitUntil(_wait_live_data_updated, timeout=7000)
-    # Check if data that is plotted is the same as the scan_item
-    # Plot may not be updated yet, so we need to wait for the data to be updated
-    qtbot.waitUntil(lambda: len(widget.curves[0].get_data()[0]) == 50)
-    scan_item = bec.queue.scan_storage.storage[-1]
-    samx_data = scan_item.live_data.samx.samx.val
-    bpm4i_data = scan_item.live_data.bpm4i.bpm4i.val
+    qtbot.waitUntil(_wait_for_scan_in_history, timeout=7000)
+
+    scan_item = bec.history.get_by_scan_id(s.scan.scan_id)
+    samx_data = scan_item.devices.samx.samx.read()["value"]
+    bpm4i_data = scan_item.devices.bpm4i.bpm4i.read()["value"]
     curve = widget.curves[0]
     assert np.allclose(curve.get_data()[0], samx_data)
     assert np.allclose(curve.get_data()[1], bpm4i_data)