test(test_bec_figure): tests for BECFigure added

bec_widgets/widgets/figure/figure.py

@@ -22,7 +22,7 @@ class FigureConfig(ConnectionConfig):
     """Configuration for BECFigure. Inheriting from ConnectionConfig widget_class and gui_id"""
 
     theme: Literal["dark", "light"] = Field("dark", description="The theme of the figure widget.")
-    num_columns: int = Field(1, description="The number of columns in the figure widget.")
+    num_cols: int = Field(1, description="The number of columns in the figure widget.")
     num_rows: int = Field(1, description="The number of rows in the figure widget.")
     widgets: dict[str, WidgetConfig] = Field(
         {}, description="The list of widgets to be added to the figure widget."
@@ -95,6 +95,8 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget):
         if config is None:
             config = FigureConfig(widget_class=self.__class__.__name__)
         else:
+            if isinstance(config, dict):
+                config = FigureConfig(**config)
             self.config = config
         super().__init__(client=client, config=config, gui_id=gui_id)
         pg.GraphicsLayoutWidget.__init__(self, parent)
@@ -160,7 +162,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget):
         if not widget_id:
             widget_id = self._generate_unique_widget_id()
         if widget_id in self.widgets:
-            raise ValueError(f"Widget with ID {widget_id} already exists.")
+            raise ValueError(f"Widget with ID '{widget_id}' already exists.")
 
         widget = self.widget_handler.create_widget(
             widget_type=widget_type,
@@ -189,10 +191,9 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget):
             # Add widget to the figure
             self.addItem(widget, row=row, col=col)
 
-        # TODO decide if needed
-        # Update num_columns and num_rows based on the added widget
+        # Update num_cols and num_rows based on the added widget
         self.config.num_rows = max(self.config.num_rows, row + 1)
-        self.config.num_columns = max(self.config.num_columns, col + 1)
+        self.config.num_cols = max(self.config.num_cols, col + 1)
 
         # Saving config for future referencing
         self.config.widgets[widget_id] = widget.config
@@ -237,12 +238,8 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget):
         widget = self._get_widget_by_coordinates(row, col)
         if widget:
             widget_id = widget.config.gui_id
-            if widget_id and widget_id in self.widgets:
+            if widget_id in self.widgets:
                 self._remove_by_id(widget_id)
-            else:
-                raise ValueError(f"No widget found at coordinates ({row}, {col}).")
-        else:
-            raise ValueError(f"No widget found at coordinates ({row}, {col}).")
 
     def _remove_by_id(self, widget_id: str) -> None:
         """
@@ -260,7 +257,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget):
                 self.config.widgets.pop(widget_id)
             print(f"Removed widget {widget_id}.")
         else:
-            raise ValueError(f"Widget with ID {widget_id} does not exist.")
+            raise ValueError(f"Widget with ID '{widget_id}' does not exist.")
 
     def __getitem__(self, key: tuple | str):
         if isinstance(key, tuple) and len(key) == 2:
@@ -287,7 +284,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget):
         """
         widget = self.getItem(row, col)
         if widget is None:
-            raise KeyError(f"No widget at coordinates ({row}, {col})")
+            raise ValueError(f"No widget at coordinates ({row}, {col})")
         return widget
 
     def _find_next_empty_position(self):
@@ -376,6 +373,9 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget):
             new_grid[row][col] = widget_id
             current_idx += 1
 
+        self.config.num_rows = row
+        self.config.num_cols = col
+
         # Update widgets' positions and replot them according to the new grid
         self.grid = new_grid
         self._reindex_grid()  # This method should be updated to handle reshuffling correctly
@@ -402,7 +402,7 @@ from qtconsole.inprocess import QtInProcessKernelManager
 from qtconsole.rich_jupyter_widget import RichJupyterWidget
 
 
-class JupyterConsoleWidget(RichJupyterWidget):
+class JupyterConsoleWidget(RichJupyterWidget):  # pragma: no cover:
     def __init__(self):
         super().__init__()
 
@@ -418,7 +418,7 @@ class JupyterConsoleWidget(RichJupyterWidget):
         self.kernel_manager.shutdown_kernel()
 
 
-class DebugWindow(QWidget):
+class DebugWindow(QWidget):  # pragma: no cover:
     """Debug window for BEC widgets"""
 
     def __init__(self, parent=None):

tests/test_bec_figure.py

@@ -0,0 +1,214 @@
+# pylint: disable = no-name-in-module,missing-class-docstring, missing-module-docstring
+import os
+
+import numpy as np
+import pytest
+from unittest.mock import MagicMock
+
+from .client_mocks import mocked_client
+from bec_widgets.widgets import BECFigure
+
+
+@pytest.fixture
+def bec_figure(qtbot, mocked_client):
+    widget = BECFigure(client=mocked_client)
+    qtbot.addWidget(widget)
+    qtbot.waitExposed(widget)
+    return widget
+
+
+def test_bec_figure_init(bec_figure):
+    assert bec_figure is not None
+    assert bec_figure.client is not None
+    assert isinstance(bec_figure, BECFigure)
+    assert bec_figure.config.widget_class == "BECFigure"
+
+
+def test_bec_figure_init_with_config(mocked_client):
+    config = {
+        "widget_class": "BECFigure",
+        "gui_id": "test_gui_id",
+        "theme": "dark",
+    }
+    widget = BECFigure(client=mocked_client, config=config)
+    assert widget.config.gui_id == "test_gui_id"
+    assert widget.config.theme == "dark"
+
+
+def test_bec_figure_add_remove_plot(bec_figure):
+    initial_count = len(bec_figure.widgets)
+
+    # Adding 3 widgets -  2 WaveformBase and 1 PlotBase
+    w0 = bec_figure.add_plot()
+    w1 = bec_figure.add_plot(widget_id="test_waveform")
+    w2 = bec_figure.add_widget(widget_id="test_plot", widget_type="PlotBase")
+
+    # Check if the widgets were added
+    assert len(bec_figure.widgets) == initial_count + 3
+    assert "widget_1" in bec_figure.widgets
+    assert "test_plot" in bec_figure.widgets
+    assert "test_waveform" in bec_figure.widgets
+    assert bec_figure.widgets["widget_1"].config.widget_class == "BECWaveform1D"
+    assert bec_figure.widgets["test_plot"].config.widget_class == "BECPlotBase"
+    assert bec_figure.widgets["test_waveform"].config.widget_class == "BECWaveform1D"
+
+    # Check accessing positions by the grid in figure
+    assert bec_figure[0, 0] == w0
+    assert bec_figure[1, 0] == w1
+    assert bec_figure[2, 0] == w2
+
+    # Removing 1 widget - PlotBase
+    bec_figure.remove(widget_id="test_plot")
+    assert len(bec_figure.widgets) == initial_count + 2
+    assert "test_plot" not in bec_figure.widgets
+    assert "test_waveform" in bec_figure.widgets
+    assert bec_figure.widgets["test_waveform"].config.widget_class == "BECWaveform1D"
+
+
+def test_access_widgets_access_errors(bec_figure):
+    bec_figure.add_plot(widget_id="test_waveform_1", row=0, col=0)
+
+    # access widget by non-existent coordinates
+    with pytest.raises(ValueError) as excinfo:
+        bec_figure[0, 2]
+        assert "No widget at coordinates (0, 2)" in str(excinfo.value)
+
+    # access widget by non-existent widget_id
+    with pytest.raises(KeyError) as excinfo:
+        bec_figure["non_existent_widget"]
+        assert "Widget with id 'non_existent_widget' not found" in str(excinfo.value)
+
+    # access widget by wrong type
+    with pytest.raises(TypeError) as excinfo:
+        bec_figure[1.2]
+        assert (
+            "Key must be a string (widget id) or a tuple of two integers (grid coordinates)"
+            in str(excinfo.value)
+        )
+
+
+def test_add_plot_to_occupied_position(bec_figure):
+    bec_figure.add_plot(widget_id="test_waveform_1", row=0, col=0)
+
+    with pytest.raises(ValueError) as excinfo:
+        bec_figure.add_plot(widget_id="test_waveform_2", row=0, col=0)
+        assert "Position at row 0 and column 0 is already occupied." in str(excinfo.value)
+
+
+def test_add_plot_to_occupied_id(bec_figure):
+    bec_figure.add_plot(widget_id="test_waveform", row=0, col=0)
+
+    with pytest.raises(ValueError) as excinfo:
+        bec_figure.add_plot(widget_id="test_waveform", row=0, col=1)
+        assert "Widget with ID 'test_waveform' already exists" in str(excinfo.value)
+
+
+def test_remove_plots(bec_figure):
+    w1 = bec_figure.add_plot(widget_id="test_waveform_1", row=0, col=0)
+    w2 = bec_figure.add_plot(widget_id="test_waveform_2", row=0, col=1)
+    w3 = bec_figure.add_plot(widget_id="test_waveform_3", row=1, col=0)
+    w4 = bec_figure.add_plot(widget_id="test_waveform_4", row=1, col=1)
+
+    assert bec_figure[0, 0] == w1
+    assert bec_figure[0, 1] == w2
+    assert bec_figure[1, 0] == w3
+    assert bec_figure[1, 1] == w4
+
+    # remove by coordinates
+    bec_figure[0, 0].remove()
+    assert "test_waveform_1" not in bec_figure.widgets
+
+    # remove by widget_id
+    bec_figure.remove(widget_id="test_waveform_2")
+    assert "test_waveform_2" not in bec_figure.widgets
+
+    # remove by widget object
+    w3.remove()
+    assert "test_waveform_3" not in bec_figure.widgets
+
+    # check the remaining widget 4
+    assert bec_figure[0, 0] == w4
+    assert bec_figure["test_waveform_4"] == w4
+    assert "test_waveform_4" in bec_figure.widgets
+    assert len(bec_figure.widgets) == 1
+
+
+def test_remove_plots_by_coordinates_ints(bec_figure):
+    w1 = bec_figure.add_plot(widget_id="test_waveform_1", row=0, col=0)
+    w2 = bec_figure.add_plot(widget_id="test_waveform_2", row=0, col=1)
+
+    bec_figure.remove(0, 0)
+    assert "test_waveform_1" not in bec_figure.widgets
+    assert "test_waveform_2" in bec_figure.widgets
+    assert bec_figure[0, 0] == w2
+    assert len(bec_figure.widgets) == 1
+
+
+def test_remove_plots_by_coordinates_tuple(bec_figure):
+    w1 = bec_figure.add_plot(widget_id="test_waveform_1", row=0, col=0)
+    w2 = bec_figure.add_plot(widget_id="test_waveform_2", row=0, col=1)
+
+    bec_figure.remove(coordinates=(0, 0))
+    assert "test_waveform_1" not in bec_figure.widgets
+    assert "test_waveform_2" in bec_figure.widgets
+    assert bec_figure[0, 0] == w2
+    assert len(bec_figure.widgets) == 1
+
+
+def test_remove_plot_by_id_error(bec_figure):
+    bec_figure.add_plot(widget_id="test_waveform_1", row=0, col=0)
+
+    with pytest.raises(ValueError) as excinfo:
+        bec_figure.remove(widget_id="non_existent_widget")
+        assert "Widget with ID 'non_existent_widget' does not exist." in str(excinfo.value)
+
+
+def test_remove_plot_by_coordinates_error(bec_figure):
+    bec_figure.add_plot(widget_id="test_waveform_1", row=0, col=0)
+
+    with pytest.raises(ValueError) as excinfo:
+        bec_figure.remove(0, 1)
+        assert "No widget at coordinates (0, 1)" in str(excinfo.value)
+
+
+def test_remove_plot_by_providing_nothing(bec_figure):
+    bec_figure.add_plot(widget_id="test_waveform_1", row=0, col=0)
+
+    with pytest.raises(ValueError) as excinfo:
+        bec_figure.remove()
+        assert "Must provide either widget_id or coordinates for removal." in str(excinfo.value)
+
+
+def test_change_theme(bec_figure):
+    bec_figure.change_theme("dark")
+    assert bec_figure.config.theme == "dark"
+    assert bec_figure.backgroundBrush().color().name() == "#000000"
+    bec_figure.change_theme("light")
+    assert bec_figure.config.theme == "light"
+    assert bec_figure.backgroundBrush().color().name() == "#ffffff"
+    bec_figure.change_theme("dark")
+    assert bec_figure.config.theme == "dark"
+    assert bec_figure.backgroundBrush().color().name() == "#000000"
+
+
+def test_change_layout(bec_figure):
+    w1 = bec_figure.add_plot(widget_id="test_waveform_1", row=0, col=0)
+    w2 = bec_figure.add_plot(widget_id="test_waveform_2", row=0, col=1)
+    w3 = bec_figure.add_plot(widget_id="test_waveform_3", row=1, col=0)
+    w4 = bec_figure.add_plot(widget_id="test_waveform_4", row=1, col=1)
+
+    bec_figure.change_layout(max_columns=1)
+
+    assert np.shape(bec_figure.grid) == (4, 1)
+    assert bec_figure[0, 0] == w1
+    assert bec_figure[1, 0] == w2
+    assert bec_figure[2, 0] == w3
+    assert bec_figure[3, 0] == w4
+
+    bec_figure.change_layout(max_rows=1)
+
+    assert np.shape(bec_figure.grid) == (1, 4)
+    assert bec_figure[0, 0] == w1
+    assert bec_figure[0, 1] == w2
+    assert bec_figure[0, 2] == w3
+    assert bec_figure[0, 3] == w4