from unittest import mock import numpy as np import pytest from bec_lib import messages from bec_lib.scan_history import ScanHistory 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={}, request_inputs={"arg_bundle": ["samx", -5, 5, 10, "samy", -5, 5, 10], "kwargs": {}}, ) heatmap_widget.plot(x_name="samx", y_name="samy", z_name="bpm4i") 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.currentText() == heatmap_widget._image_config.x_device.name dialog.reject() qtbot.waitUntil(lambda: heatmap_widget.heatmap_dialog is None) def test_heatmap_widget_reset(heatmap_widget): """ Test that the reset method clears the plot. """ heatmap_widget.scan_item = create_dummy_scan_item() heatmap_widget.plot(x_name="samx", y_name="samy", z_name="bpm4i") heatmap_widget.reset() assert heatmap_widget._grid_index is None assert heatmap_widget.main_image.raw_data is None def test_heatmap_widget_update_plot_with_scan_history(heatmap_widget, grid_scan_history_msg, qtbot): """ Test that the update_plot method updates the plot with scan history. """ heatmap_widget.client.history = ScanHistory(heatmap_widget.client, False) heatmap_widget.client.history._scan_data[grid_scan_history_msg.scan_id] = grid_scan_history_msg heatmap_widget.client.history._scan_ids.append(grid_scan_history_msg.scan_id) heatmap_widget.client.queue.scan_storage.current_scan = None heatmap_widget.plot( x_name="samx", y_name="samy", z_name="bpm4i", x_entry="samx", y_entry="samy", z_entry="bpm4i", ) qtbot.waitUntil(lambda: heatmap_widget.main_image.raw_data is not None) qtbot.waitUntil(lambda: heatmap_widget.main_image.raw_data.shape == (10, 10)) heatmap_widget.enforce_interpolation = True heatmap_widget.oversampling_factor = 2.0 qtbot.waitUntil(lambda: heatmap_widget.main_image.raw_data.shape == (20, 20))