From b1aff6d791ff847eb2f628e66ccaa4672fdeea08 Mon Sep 17 00:00:00 2001
From: appel_c <christian.appel@psi.ch>
Date: Fri, 6 Sep 2024 12:45:47 +0200
Subject: [PATCH] test: add tests, including extension to end-2-end test

---
 tests/end-2-end/test_bec_figure_rpc_e2e.py |  7 +++
 tests/unit_tests/test_waveform_widget.py   | 61 ++++++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/tests/end-2-end/test_bec_figure_rpc_e2e.py b/tests/end-2-end/test_bec_figure_rpc_e2e.py
index 0b1268db..24f06a22 100644
--- a/tests/end-2-end/test_bec_figure_rpc_e2e.py
+++ b/tests/end-2-end/test_bec_figure_rpc_e2e.py
@@ -188,6 +188,13 @@ def test_dap_rpc(rpc_server_figure, bec_client_lib, qtbot):
 
     qtbot.waitUntil(wait_for_fit, timeout=10000)
 
+    # Repeat fit after adding a region of interest
+    plt.select_roi(region=(3, 7))
+    res = scans.line_scan(dev.samx, 0, 8, steps=50, relative=False)
+    res.wait()
+
+    qtbot.waitUntil(wait_for_fit, timeout=10000)
+
 
 def test_removing_subplots(rpc_server_figure, bec_client_lib):
     fig = BECFigure(rpc_server_figure)
diff --git a/tests/unit_tests/test_waveform_widget.py b/tests/unit_tests/test_waveform_widget.py
index a6824151..0bf07e90 100644
--- a/tests/unit_tests/test_waveform_widget.py
+++ b/tests/unit_tests/test_waveform_widget.py
@@ -7,6 +7,7 @@ from qtpy.QtWidgets import QApplication
 
 from bec_widgets.qt_utils.settings_dialog import SettingsDialog
 from bec_widgets.utils.colors import apply_theme, get_theme_palette, set_theme
+from bec_widgets.utils.linear_region_selector import LinearRegionWrapper
 from bec_widgets.widgets.figure.plots.axis_settings import AxisSettings
 from bec_widgets.widgets.waveform.waveform_popups.curve_dialog.curve_dialog import CurveSettings
 from bec_widgets.widgets.waveform.waveform_popups.dap_summary_dialog.dap_summary_dialog import (
@@ -518,3 +519,63 @@ def test_waveform_widget_theme_update(qtbot, waveform_widget):
     bg_color = waveform_widget.fig.backgroundBrush().color()
     assert bg_color == QColor("black")
     assert waveform_color == waveform_color_dark
+
+
+def test_waveform_roi_selection_creation(waveform_widget, qtbot):
+    """Test ROI selection for waveform widget.
+
+    This checks that the ROI select is properly created and removed when the button is toggled.
+    """
+    # Check if curve is create upon ROI select slot
+    # This also checks that the button in the toolbar works
+    container = []
+
+    def callback(msg):
+        container.append(msg)
+
+    waveform_widget.waveform.roi_active.connect(callback)
+    assert waveform_widget.waveform.roi_select is None
+    assert waveform_widget.waveform.roi_region == (None, None)
+    # Toggle the ROI select
+    waveform_widget.toogle_roi_select(True)
+    assert isinstance(waveform_widget.waveform.roi_select, LinearRegionWrapper)
+    # This is the default region for the pg.LinearRegionItem
+    assert waveform_widget.waveform.roi_region == (0, 1)
+    # Untoggle the ROI select
+    waveform_widget.toogle_roi_select(False)
+    assert waveform_widget.waveform.roi_select is None
+    assert container[0] is True
+    assert container[1] is False
+
+
+def test_waveform_roi_selection_updates_fit(waveform_widget, qtbot):
+    """This test checks that upon selection of a new region, the fit is updated and all signals are emitted as expected."""
+    container = []
+
+    def callback(msg):
+        container.append(msg)
+
+    waveform_widget.waveform.roi_changed.connect(callback)
+    # Mock refresh_dap method
+    with patch.object(waveform_widget.waveform, "refresh_dap") as mock_refresh_dap:
+        waveform_widget.toogle_roi_select(True)
+        waveform_widget.waveform.roi_select.linear_region_selector.setRegion([0.5, 1.5])
+        qtbot.wait(200)
+        assert waveform_widget.waveform.roi_region == (0.5, 1.5)
+        waveform_widget.toogle_roi_select(False)
+        assert waveform_widget.waveform.roi_region == (None, None)
+        assert len(container) == 1
+        assert container[0] == (0.5, 1.5)
+        # 3 refresh DAP calls: 1x upon hook, 1x unhook and 1x from roi_changed
+        assert mock_refresh_dap.call_count == 3
+
+
+def test_waveform_roi_selection_change_color(waveform_widget, qtbot):
+    """This test checks that the color of the ROI region can be changed."""
+    waveform_widget.toogle_roi_select(True)
+    waveform_widget.waveform.roi_select.change_roi_color((QColor("red"), QColor("blue")))
+    # I can only get the brush from the RegionSelectItem
+    assert (
+        waveform_widget.waveform.roi_select.linear_region_selector.currentBrush.color()
+        == QColor("red")
+    )