feat: add heatmap widget

tests/unit_tests/test_heatmap_widget.py

@@ -0,0 +1,322 @@
+from unittest import mock
+
+import numpy as np
+import pytest
+from bec_lib import messages
+
+from bec_widgets.widgets.plots.heatmap.heatmap import Heatmap, HeatmapConfig, HeatmapDeviceSignal
+
+# pytest: disable=unused-import
+from tests.unit_tests.client_mocks import mocked_client
+
+from .client_mocks import create_dummy_scan_item
+
+
+@pytest.fixture
+def heatmap_widget(qtbot, mocked_client):
+    widget = Heatmap(client=mocked_client)
+    qtbot.addWidget(widget)
+    qtbot.waitExposed(widget)
+    yield widget
+
+
+def test_heatmap_plot(heatmap_widget):
+    heatmap_widget.plot(x_name="samx", y_name="samy", z_name="bpm4i")
+
+    assert heatmap_widget._image_config.x_device.name == "samx"
+    assert heatmap_widget._image_config.y_device.name == "samy"
+    assert heatmap_widget._image_config.z_device.name == "bpm4i"
+
+
+def test_heatmap_on_scan_status_no_scan_id(heatmap_widget):
+
+    scan_msg = messages.ScanStatusMessage(scan_id=None, status="open", metadata={}, info={})
+    with mock.patch.object(heatmap_widget, "reset") as mock_reset:
+
+        heatmap_widget.on_scan_status(scan_msg.content, scan_msg.metadata)
+        mock_reset.assert_not_called()
+
+
+def test_heatmap_on_scan_status_same_scan_id(heatmap_widget):
+    scan_msg = messages.ScanStatusMessage(scan_id="123", status="open", metadata={}, info={})
+    heatmap_widget.scan_id = "123"
+    with mock.patch.object(heatmap_widget, "reset") as mock_reset:
+        heatmap_widget.on_scan_status(scan_msg.content, scan_msg.metadata)
+        mock_reset.assert_not_called()
+
+
+def test_heatmap_widget_on_scan_status_different_scan_id(heatmap_widget):
+    scan_msg = messages.ScanStatusMessage(scan_id="123", status="open", metadata={}, info={})
+    heatmap_widget.scan_id = "456"
+    with mock.patch.object(heatmap_widget, "reset") as mock_reset:
+        heatmap_widget.on_scan_status(scan_msg.content, scan_msg.metadata)
+        mock_reset.assert_called_once()
+
+
+def test_heatmap_get_image_data_missing_data(heatmap_widget):
+    """
+    If the data is missing or incomplete, the method should return None.
+    """
+    assert heatmap_widget.get_image_data() == (None, None)
+
+
+def test_heatmap_get_image_data_grid_scan(heatmap_widget):
+    scan_msg = messages.ScanStatusMessage(
+        scan_id="123", status="open", scan_name="grid_scan", metadata={}, info={}
+    )
+    heatmap_widget.status_message = scan_msg
+    with mock.patch.object(heatmap_widget, "get_grid_scan_image") as mock_get_grid_scan_image:
+        heatmap_widget.get_image_data(x_data=[1, 2], y_data=[3, 4], z_data=[5, 6])
+        mock_get_grid_scan_image.assert_called_once()
+
+
+def test_heatmap_get_image_data_step_scan(heatmap_widget):
+    """
+    If the step scan has too few points, it should return None.
+    """
+    scan_msg = messages.ScanStatusMessage(
+        scan_id="123",
+        status="open",
+        scan_name="step_scan",
+        scan_type="step",
+        metadata={},
+        info={"positions": [[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]},
+    )
+    with mock.patch.object(heatmap_widget, "get_step_scan_image") as mock_get_step_scan_image:
+        heatmap_widget.status_message = scan_msg
+        heatmap_widget.get_image_data(x_data=[1, 2, 3, 4], y_data=[1, 2, 3, 4], z_data=[1, 2, 5, 6])
+        mock_get_step_scan_image.assert_called_once()
+
+
+def test_heatmap_get_image_data_step_scan_too_few_points(heatmap_widget):
+    """
+    If the step scan has too few points, it should return None.
+    """
+    scan_msg = messages.ScanStatusMessage(
+        scan_id="123",
+        status="open",
+        scan_name="step_scan",
+        scan_type="step",
+        metadata={},
+        info={"positions": [[1, 2], [3, 4]]},
+    )
+    heatmap_widget.status_message = scan_msg
+    out = heatmap_widget.get_image_data(x_data=[1, 2], y_data=[3, 4], z_data=[5, 6])
+    assert out == (None, None)
+
+
+def test_heatmap_get_image_data_unsupported_scan(heatmap_widget):
+    scan_msg = messages.ScanStatusMessage(
+        scan_id="123", status="open", scan_type="fly", metadata={}, info={}
+    )
+    heatmap_widget.status_message = scan_msg
+    assert heatmap_widget.get_image_data(x_data=[1, 2], y_data=[3, 4], z_data=[5, 6]) == (
+        None,
+        None,
+    )
+
+
+def test_heatmap_get_grid_scan_image(heatmap_widget):
+    scan_msg = messages.ScanStatusMessage(
+        scan_id="123",
+        status="open",
+        scan_name="grid_scan",
+        metadata={},
+        info={"positions": np.random.rand(100, 2).tolist()},
+        request_inputs={"arg_bundle": ["samx", -5, 5, 10, "samy", -5, 5, 10], "kwargs": {}},
+    )
+    heatmap_widget._image_config = HeatmapConfig(
+        parent_id="parent_id",
+        x_device=HeatmapDeviceSignal(name="samx", entry="samx"),
+        y_device=HeatmapDeviceSignal(name="samy", entry="samy"),
+        z_device=HeatmapDeviceSignal(name="bpm4i", entry="bpm4i"),
+        color_map="viridis",
+    )
+    img, _ = heatmap_widget.get_grid_scan_image(list(range(100)), msg=scan_msg)
+    assert img.shape == (10, 10)
+    assert sorted(np.asarray(img, dtype=int).flatten().tolist()) == list(range(100))
+
+
+def test_heatmap_get_step_scan_image(heatmap_widget):
+
+    scan_msg = messages.ScanStatusMessage(
+        scan_id="123",
+        status="open",
+        scan_name="step_scan",
+        scan_type="step",
+        metadata={},
+        info={"positions": np.random.rand(100, 2).tolist()},
+    )
+    heatmap_widget.status_message = scan_msg
+    heatmap_widget.scan_item = create_dummy_scan_item()
+    heatmap_widget.scan_item.status_message = scan_msg
+    heatmap_widget._image_config = HeatmapConfig(
+        parent_id="parent_id",
+        x_device=HeatmapDeviceSignal(name="samx", entry="samx"),
+        y_device=HeatmapDeviceSignal(name="samy", entry="samy"),
+        z_device=HeatmapDeviceSignal(name="bpm4i", entry="bpm4i"),
+        color_map="viridis",
+    )
+    img, _ = heatmap_widget.get_step_scan_image(
+        list(np.random.rand(100)), list(np.random.rand(100)), list(range(100)), msg=scan_msg
+    )
+    assert img.shape > (10, 10)
+
+
+def test_heatmap_update_plot_no_scan_item(heatmap_widget):
+    heatmap_widget._image_config = HeatmapConfig(
+        parent_id="parent_id",
+        x_device=HeatmapDeviceSignal(name="samx", entry="samx"),
+        y_device=HeatmapDeviceSignal(name="samy", entry="samy"),
+        z_device=HeatmapDeviceSignal(name="bpm4i", entry="bpm4i"),
+        color_map="viridis",
+    )
+    with mock.patch.object(heatmap_widget.main_image, "setImage") as mock_set_image:
+        heatmap_widget.update_plot(_override_slot_params={"verify_sender": False})
+        mock_set_image.assert_not_called()
+
+
+def test_heatmap_update_plot(heatmap_widget):
+    heatmap_widget._image_config = HeatmapConfig(
+        parent_id="parent_id",
+        x_device=HeatmapDeviceSignal(name="samx", entry="samx"),
+        y_device=HeatmapDeviceSignal(name="samy", entry="samy"),
+        z_device=HeatmapDeviceSignal(name="bpm4i", entry="bpm4i"),
+        color_map="viridis",
+    )
+    heatmap_widget.scan_item = create_dummy_scan_item()
+    heatmap_widget.scan_item.status_message = messages.ScanStatusMessage(
+        scan_id="123",
+        status="open",
+        scan_name="grid_scan",
+        metadata={},
+        info={"positions": np.random.rand(100, 2).tolist()},
+        request_inputs={"arg_bundle": ["samx", -5, 5, 10, "samy", -5, 5, 10], "kwargs": {}},
+    )
+    with mock.patch.object(heatmap_widget.main_image, "setImage") as mock_set_image:
+        heatmap_widget.update_plot(_override_slot_params={"verify_sender": False})
+        img = mock_set_image.mock_calls[0].args[0]
+        assert img.shape == (10, 10)
+
+
+def test_heatmap_update_plot_without_status_message(heatmap_widget):
+    heatmap_widget._image_config = HeatmapConfig(
+        parent_id="parent_id",
+        x_device=HeatmapDeviceSignal(name="samx", entry="samx"),
+        y_device=HeatmapDeviceSignal(name="samy", entry="samy"),
+        z_device=HeatmapDeviceSignal(name="bpm4i", entry="bpm4i"),
+        color_map="viridis",
+    )
+    heatmap_widget.scan_item = create_dummy_scan_item()
+    heatmap_widget.scan_item.status_message = None
+    with mock.patch.object(heatmap_widget.main_image, "setImage") as mock_set_image:
+        heatmap_widget.update_plot(_override_slot_params={"verify_sender": False})
+        mock_set_image.assert_not_called()
+
+
+def test_heatmap_update_plot_no_img_data(heatmap_widget):
+    heatmap_widget._image_config = HeatmapConfig(
+        parent_id="parent_id",
+        x_device=HeatmapDeviceSignal(name="samx", entry="samx"),
+        y_device=HeatmapDeviceSignal(name="samy", entry="samy"),
+        z_device=HeatmapDeviceSignal(name="bpm4i", entry="bpm4i"),
+        color_map="viridis",
+    )
+    heatmap_widget.scan_item = create_dummy_scan_item()
+    heatmap_widget.scan_item.status_message = messages.ScanStatusMessage(
+        scan_id="123",
+        status="open",
+        scan_name="grid_scan",
+        metadata={},
+        info={},
+        request_inputs={"arg_bundle": ["samx", -5, 5, 10, "samy", -5, 5, 10], "kwargs": {}},
+    )
+    with mock.patch.object(heatmap_widget, "get_image_data", return_value=None):
+        with mock.patch.object(heatmap_widget.main_image, "setImage") as mock_set_image:
+            heatmap_widget.update_plot(_override_slot_params={"verify_sender": False})
+            mock_set_image.assert_not_called()
+
+
+def test_heatmap_settings_popup(heatmap_widget, qtbot):
+    """
+    Test that the settings popup opens and contains the expected elements.
+    """
+    settings_action = heatmap_widget.toolbar.components.get_action("heatmap_settings").action
+    heatmap_widget.show_heatmap_settings()
+    qtbot.waitUntil(lambda: heatmap_widget.heatmap_dialog is not None)
+
+    assert heatmap_widget.heatmap_dialog.isVisible()
+
+    assert settings_action.isChecked()
+
+    heatmap_widget.heatmap_dialog.reject()
+    qtbot.waitUntil(lambda: heatmap_widget.heatmap_dialog is None)
+
+    assert not settings_action.isChecked()
+
+
+def test_heatmap_settings_popup_already_open(heatmap_widget, qtbot):
+    """
+    Test that if the settings dialog is already open, it is brought to the front.
+    """
+    heatmap_widget.show_heatmap_settings()
+    qtbot.waitUntil(lambda: heatmap_widget.heatmap_dialog is not None)
+
+    initial_dialog = heatmap_widget.heatmap_dialog
+
+    heatmap_widget.show_heatmap_settings()
+    qtbot.waitUntil(lambda: heatmap_widget.heatmap_dialog is initial_dialog)
+
+    assert heatmap_widget.heatmap_dialog.isVisible()  # Dialog should still be visible
+    assert heatmap_widget.heatmap_dialog is initial_dialog  # Should be the same dialog
+
+    heatmap_widget.heatmap_dialog.reject()
+    qtbot.waitUntil(lambda: heatmap_widget.heatmap_dialog is None)
+
+
+def test_heatmap_settings_popup_accept_changes(heatmap_widget, qtbot):
+    """
+    Test that changes made in the settings dialog are applied correctly.
+    """
+    heatmap_widget.plot(x_name="samx", y_name="samy", z_name="bpm4i")
+    assert heatmap_widget.color_map == "plasma"  # Default colormap
+    heatmap_widget.show_heatmap_settings()
+    qtbot.waitUntil(lambda: heatmap_widget.heatmap_dialog is not None)
+
+    dialog = heatmap_widget.heatmap_dialog
+    assert dialog.widget.isVisible()
+
+    # Simulate changing a setting
+    dialog.widget.ui.color_map.colormap = "viridis"
+
+    # Accept changes
+    dialog.accept()
+
+    qtbot.waitUntil(lambda: heatmap_widget.heatmap_dialog is None)
+
+    # Verify that the setting was applied
+    assert heatmap_widget.color_map == "viridis"
+
+
+def test_heatmap_settings_popup_show_settings(heatmap_widget, qtbot):
+    """
+    Test that the settings dialog opens and contains the expected elements.
+    """
+    heatmap_widget.plot(x_name="samx", y_name="samy", z_name="bpm4i")
+    heatmap_widget.show_heatmap_settings()
+    qtbot.waitUntil(lambda: heatmap_widget.heatmap_dialog is not None)
+
+    dialog = heatmap_widget.heatmap_dialog
+    assert dialog.isVisible()
+    assert dialog.widget is not None
+    assert hasattr(dialog.widget.ui, "color_map")
+    assert hasattr(dialog.widget.ui, "x_name")
+    assert hasattr(dialog.widget.ui, "y_name")
+    assert hasattr(dialog.widget.ui, "z_name")
+
+    # Check that the ui elements are correctly initialized
+    assert dialog.widget.ui.color_map.colormap == heatmap_widget.color_map
+    assert dialog.widget.ui.x_name.text() == heatmap_widget._image_config.x_device.name
+
+    dialog.reject()
+    qtbot.waitUntil(lambda: heatmap_widget.heatmap_dialog is None)