feat(waveform): new Waveform widget based on NextGen PlotBase

tests/unit_tests/test_curve_settings.py

@@ -0,0 +1,367 @@
+import json
+from unittest.mock import MagicMock, patch
+
+import pytest
+from qtpy.QtWidgets import QComboBox, QVBoxLayout
+
+from bec_widgets.widgets.plots_next_gen.waveform.settings.curve_settings.curve_setting import (
+    CurveSetting,
+)
+from bec_widgets.widgets.plots_next_gen.waveform.settings.curve_settings.curve_tree import CurveTree
+from bec_widgets.widgets.plots_next_gen.waveform.waveform import Waveform
+from tests.unit_tests.client_mocks import dap_plugin_message, mocked_client, mocked_client_with_dap
+from tests.unit_tests.conftest import create_widget
+
+##################################################
+# CurveSetting
+##################################################
+
+
+@pytest.fixture
+def curve_setting_fixture(qtbot, mocked_client):
+    """
+    Creates a CurveSetting widget targeting a mock or real Waveform widget.
+    """
+    wf = create_widget(qtbot, Waveform, client=mocked_client)
+    wf.x_mode = "auto"
+    curve_setting = create_widget(qtbot, CurveSetting, parent=None, target_widget=wf)
+    return curve_setting, wf
+
+
+def test_curve_setting_init(curve_setting_fixture):
+    """
+    Ensure CurveSetting constructs properly, with a CurveTree inside
+    and an x-axis group box for modes.
+    """
+    curve_setting, wf = curve_setting_fixture
+
+    # Basic checks
+    assert curve_setting.objectName() == "CurveSetting"
+    # The layout should be QVBoxLayout
+    assert isinstance(curve_setting.layout, QVBoxLayout)
+
+    # There's an x_axis_box group and a y_axis_box group
+    assert hasattr(curve_setting, "x_axis_box")
+    assert hasattr(curve_setting, "y_axis_box")
+
+    # The x_axis_box should contain a QComboBox for mode
+    mode_combo = curve_setting.mode_combo
+    assert isinstance(mode_combo, QComboBox)
+    # Should contain these items: ["auto", "index", "timestamp", "device"]
+    expected_modes = ["auto", "index", "timestamp", "device"]
+    for m in expected_modes:
+        assert m in [
+            curve_setting.mode_combo.itemText(i) for i in range(curve_setting.mode_combo.count())
+        ]
+
+    # Check that there's a curve_manager inside y_axis_box
+    assert hasattr(curve_setting, "curve_manager")
+    assert curve_setting.y_axis_box.layout.count() > 0
+
+
+def test_curve_setting_accept_changes(curve_setting_fixture, qtbot):
+    """
+    Test that calling accept_changes() applies x-axis mode changes
+    and triggers the CurveTree to send its curve JSON to the target waveform.
+    """
+    curve_setting, wf = curve_setting_fixture
+
+    # Suppose user chooses "index" from the combo
+    curve_setting.mode_combo.setCurrentText("index")
+    # The device_x is disabled if not device mode
+
+    # Spy on 'send_curve_json' from the curve_manager
+    send_spy = MagicMock()
+    curve_setting.curve_manager.send_curve_json = send_spy
+
+    # Call accept_changes()
+    curve_setting.accept_changes()
+
+    # Check that we updated the waveform
+    assert wf.x_mode == "index"
+    # Check that the manager send_curve_json was called
+    send_spy.assert_called_once()
+
+
+def test_curve_setting_switch_device_mode(curve_setting_fixture, qtbot):
+    """
+    If user chooses device mode from the combo, the device_x line edit should be enabled
+    and set to the current wavefrom.x_axis_mode["name"].
+    """
+    curve_setting, wf = curve_setting_fixture
+
+    # Initially we assume "auto"
+    assert curve_setting.mode_combo.currentText() == "auto"
+    # Switch to device
+    curve_setting.mode_combo.setCurrentText("device")
+    assert curve_setting.device_x.isEnabled()
+
+    # This line edit should reflect the waveform.x_axis_mode["name"], or be blank if none
+    assert curve_setting.device_x.text() == wf.x_axis_mode["name"]
+
+
+def test_curve_setting_refresh(curve_setting_fixture, qtbot):
+    """
+    Test that calling refresh() refreshes the embedded CurveTree
+    and re-reads the x axis mode from the waveform.
+    """
+    curve_setting, wf = curve_setting_fixture
+
+    # Suppose the waveform changed x_mode from "auto" to "timestamp" behind the scenes
+    wf.x_mode = "timestamp"
+    # Spy on the curve_manager
+    refresh_spy = MagicMock()
+    curve_setting.curve_manager.refresh_from_waveform = refresh_spy
+
+    # Call refresh
+    curve_setting.refresh()
+
+    refresh_spy.assert_called_once()
+    # The combo should now read "timestamp"
+    assert curve_setting.mode_combo.currentText() == "timestamp"
+
+
+##################################################
+# CurveTree
+##################################################
+
+
+@pytest.fixture
+def curve_tree_fixture(qtbot, mocked_client_with_dap):
+    """
+    Creates a CurveTree widget referencing a mocked or real Waveform.
+    """
+    wf = create_widget(qtbot, Waveform, client=mocked_client_with_dap)
+    wf.color_palette = "magma"
+    curve_tree = create_widget(qtbot, CurveTree, parent=None, waveform=wf)
+    return curve_tree, wf
+
+
+def test_curve_tree_init(curve_tree_fixture):
+    """
+    Test that the CurveTree initializes properly with references to the waveform,
+    sets up the toolbar, and an empty QTreeWidget.
+    """
+    curve_tree, wf = curve_tree_fixture
+    assert curve_tree.waveform == wf
+    assert curve_tree.color_palette == "magma"
+    assert curve_tree.tree.columnCount() == 7
+
+    assert "add" in curve_tree.toolbar.widgets
+    assert "expand_all" in curve_tree.toolbar.widgets
+    assert "collapse_all" in curve_tree.toolbar.widgets
+    assert "renormalize_colors" in curve_tree.toolbar.widgets
+
+
+def test_add_new_curve(curve_tree_fixture):
+    """
+    Test that add_new_curve() adds a top-level item with a device curve config,
+    assigns it a color from the buffer, and doesn't modify existing rows.
+    """
+    curve_tree, wf = curve_tree_fixture
+    curve_tree.color_buffer = ["#111111", "#222222", "#333333", "#444444", "#555555"]
+
+    assert curve_tree.tree.topLevelItemCount() == 0
+
+    with patch.object(curve_tree, "_ensure_color_buffer_size") as ensure_spy:
+        new_item = curve_tree.add_new_curve(name="bpm4i", entry="bpm4i")
+        ensure_spy.assert_called_once()
+
+    assert curve_tree.tree.topLevelItemCount() == 1
+    last_item = curve_tree.all_items[-1]
+    assert last_item is new_item
+    assert new_item.config.source == "device"
+    assert new_item.config.signal.name == "bpm4i"
+    assert new_item.config.signal.entry == "bpm4i"
+    assert new_item.config.color in curve_tree.color_buffer
+
+
+def test_renormalize_colors(curve_tree_fixture):
+    """
+    Test that renormalize_colors overwrites colors for all items in creation order.
+    """
+    curve_tree, wf = curve_tree_fixture
+    # Add multiple curves
+    c1 = curve_tree.add_new_curve(name="bpm4i", entry="bpm4i")
+    c2 = curve_tree.add_new_curve(name="bpm3a", entry="bpm3a")
+    curve_tree.color_buffer = []
+
+    set_color_spy_c1 = patch.object(c1.color_button, "set_color")
+    set_color_spy_c2 = patch.object(c2.color_button, "set_color")
+
+    with set_color_spy_c1 as spy1, set_color_spy_c2 as spy2:
+        curve_tree.renormalize_colors()
+        spy1.assert_called_once()
+        spy2.assert_called_once()
+
+
+def test_expand_collapse(curve_tree_fixture):
+    """
+    Test expand_all_daps() and collapse_all_daps() calls expand/collapse on every top-level item.
+    """
+    curve_tree, wf = curve_tree_fixture
+    c1 = curve_tree.add_new_curve(name="bpm4i", entry="bpm4i")
+    curve_tree.tree.expandAll()
+    expand_spy = patch.object(curve_tree.tree, "expandItem")
+    collapse_spy = patch.object(curve_tree.tree, "collapseItem")
+
+    with expand_spy as e_spy:
+        curve_tree.expand_all_daps()
+        e_spy.assert_called_once_with(c1)
+
+    with collapse_spy as c_spy:
+        curve_tree.collapse_all_daps()
+        c_spy.assert_called_once_with(c1)
+
+
+def test_send_curve_json(curve_tree_fixture, monkeypatch):
+    """
+    Test that send_curve_json sets the waveform's color_palette and curve_json
+    to the exported config from the tree.
+    """
+    curve_tree, wf = curve_tree_fixture
+    # Add multiple curves
+    curve_tree.add_new_curve(name="bpm4i", entry="bpm4i")
+    curve_tree.add_new_curve(name="bpm3a", entry="bpm3a")
+
+    curve_tree.color_palette = "viridis"
+    curve_tree.send_curve_json()
+
+    assert wf.color_palette == "viridis"
+    data = json.loads(wf.curve_json)
+    assert len(data) == 2
+    labels = [d["label"] for d in data]
+    assert "bpm4i-bpm4i" in labels
+    assert "bpm3a-bpm3a" in labels
+
+
+def test_refresh_from_waveform(qtbot, mocked_client_with_dap, monkeypatch):
+    """
+    Test that refresh_from_waveform() rebuilds the tree from the waveform's curve_json
+    """
+    patched_models = {"GaussianModel": {}, "LorentzModel": {}, "SineModel": {}}
+    monkeypatch.setattr(mocked_client_with_dap.dap, "_available_dap_plugins", patched_models)
+
+    wf = create_widget(qtbot, Waveform, client=mocked_client_with_dap)
+    wf.x_mode = "auto"
+    curve_tree = create_widget(qtbot, CurveTree, parent=None, waveform=wf)
+
+    wf.plot(arg1="bpm4i", dap="GaussianModel")
+    wf.plot(arg1="bpm3a", dap="GaussianModel")
+
+    # Clear the tree to simulate a fresh rebuild.
+    curve_tree.tree.clear()
+    curve_tree.all_items.clear()
+    assert curve_tree.tree.topLevelItemCount() == 0
+
+    # For DAP rows
+    curve_tree.refresh_from_waveform()
+    assert curve_tree.tree.topLevelItemCount() == 2
+
+
+def test_add_dap_row(curve_tree_fixture):
+    """
+    Test that add_dap_row creates a new DAP curve as a child of a device curve,
+    with the correct configuration and parent-child relationship.
+    """
+    curve_tree, wf = curve_tree_fixture
+
+    # Add a device curve first
+    device_row = curve_tree.add_new_curve(name="bpm4i", entry="bpm4i")
+    assert device_row.source == "device"
+    assert curve_tree.tree.topLevelItemCount() == 1
+    assert device_row.childCount() == 0
+
+    # Now add a DAP row to it
+    device_row.add_dap_row()
+
+    # Check that child was added
+    assert device_row.childCount() == 1
+    dap_child = device_row.child(0)
+
+    # Verify the DAP child has the correct configuration
+    assert dap_child.source == "dap"
+    assert dap_child.config.parent_label == device_row.config.label
+
+    # Check that the DAP inherits device name/entry from parent
+    assert dap_child.config.signal.name == "bpm4i"
+    assert dap_child.config.signal.entry == "bpm4i"
+
+    # Check that the item is in the curve_tree's all_items list
+    assert dap_child in curve_tree.all_items
+
+
+def test_remove_self_top_level(curve_tree_fixture):
+    """
+    Test that remove_self removes a top-level device row from the tree.
+    """
+    curve_tree, wf = curve_tree_fixture
+
+    # Add two device curves
+    row1 = curve_tree.add_new_curve(name="bpm4i", entry="bpm4i")
+    row2 = curve_tree.add_new_curve(name="bpm3a", entry="bpm3a")
+    assert curve_tree.tree.topLevelItemCount() == 2
+    assert len(curve_tree.all_items) == 2
+
+    # Remove the first row
+    row1.remove_self()
+
+    # Check that only one row remains and it's the correct one
+    assert curve_tree.tree.topLevelItemCount() == 1
+    assert curve_tree.tree.topLevelItem(0) == row2
+    assert len(curve_tree.all_items) == 1
+    assert curve_tree.all_items[0] == row2
+
+
+def test_remove_self_child(curve_tree_fixture):
+    """
+    Test that remove_self removes a child DAP row while preserving the parent device row.
+    """
+    curve_tree, wf = curve_tree_fixture
+
+    # Add a device curve and a DAP child
+    device_row = curve_tree.add_new_curve(name="bpm4i", entry="bpm4i")
+    device_row.add_dap_row()
+    dap_child = device_row.child(0)
+
+    assert curve_tree.tree.topLevelItemCount() == 1
+    assert device_row.childCount() == 1
+    assert len(curve_tree.all_items) == 2
+
+    # Remove the DAP child
+    dap_child.remove_self()
+
+    # Check that the parent device row still exists but has no children
+    assert curve_tree.tree.topLevelItemCount() == 1
+    assert device_row.childCount() == 0
+    assert len(curve_tree.all_items) == 1
+    assert curve_tree.all_items[0] == device_row
+
+
+def test_export_data_dap(curve_tree_fixture):
+    """
+    Test that export_data from a DAP row correctly includes parent relationship and DAP model.
+    """
+    curve_tree, wf = curve_tree_fixture
+
+    # Add a device curve with specific parameters
+    device_row = curve_tree.add_new_curve(name="bpm4i", entry="bpm4i")
+    device_row.config.label = "bpm4i-main"
+
+    # Add a DAP child
+    device_row.add_dap_row()
+    dap_child = device_row.child(0)
+
+    # Set a specific model in the DAP combobox
+    dap_child.dap_combo.fit_model_combobox.setCurrentText("GaussianModel")
+
+    # Export data from the DAP row
+    exported = dap_child.export_data()
+
+    # Check the exported data
+    assert exported["source"] == "dap"
+    assert exported["parent_label"] == "bpm4i-main"
+    assert exported["signal"]["name"] == "bpm4i"
+    assert exported["signal"]["entry"] == "bpm4i"
+    assert exported["signal"]["dap"] == "GaussianModel"
+    assert exported["label"] == "bpm4i-main-GaussianModel"