From 46ee64f29dd5552ae92a7047c29faf2f37c480f9 Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Mon, 12 Jun 2023 17:37:33 +0100
Subject: [PATCH] Test layer changes for ScatterWidget

---
 .../tests/test_layer_changes.py               | 64 +++++++++++--------
 1 file changed, 36 insertions(+), 28 deletions(-)

diff --git a/src/napari_matplotlib/tests/test_layer_changes.py b/src/napari_matplotlib/tests/test_layer_changes.py
index 87baebfd..bdd6c600 100644
--- a/src/napari_matplotlib/tests/test_layer_changes.py
+++ b/src/napari_matplotlib/tests/test_layer_changes.py
@@ -9,6 +9,7 @@
 from napari_matplotlib import (
     FeaturesScatterWidget,
     HistogramWidget,
+    ScatterWidget,
     SliceWidget,
 )
 from napari_matplotlib.base import NapariMPLWidget
@@ -18,34 +19,32 @@
 )
 
 
-@pytest.mark.parametrize("widget_cls", [HistogramWidget, SliceWidget])
+@pytest.mark.parametrize(
+    "widget_cls, n_layers",
+    [(HistogramWidget, 1), (SliceWidget, 1), (ScatterWidget, 2)],
+)
 def test_change_one_layer(
-    make_napari_viewer, brain_data, astronaut_data, widget_cls
+    make_napari_viewer,
+    brain_data,
+    astronaut_data,
+    widget_cls,
+    n_layers,
 ):
     """
     Test all widgets that take one layer as input to make sure the plot changes
     when the napari layer selection changes.
     """
     viewer = make_napari_viewer()
-    assert_one_layer_plot_changes(
-        viewer, widget_cls, brain_data, astronaut_data
-    )
-
 
-def assert_one_layer_plot_changes(
-    viewer: Viewer,
-    widget_cls: Type[NapariMPLWidget],
-    data1: Tuple[npt.NDArray[np.generic], Dict[str, Any]],
-    data2: Tuple[npt.NDArray[np.generic], Dict[str, Any]],
-) -> None:
-    """
-    When the selected layer is changed, make sure the plot generated
-    by `widget_cls` also changes.
-    """
     widget = widget_cls(viewer)
-    viewer.add_image(data1[0], **data1[1])
-    viewer.add_image(data2[0], **data2[1])
-    assert_plot_changes(viewer, widget)
+    # Add n copies of two different datasets
+    for _ in range(n_layers):
+        viewer.add_image(brain_data[0], **brain_data[1])
+    for _ in range(n_layers):
+        viewer.add_image(astronaut_data[0], **astronaut_data[1])
+
+    assert len(viewer.layers) == 2 * n_layers
+    assert_plot_changes(viewer, widget, n_layers=n_layers)
 
 
 @pytest.mark.parametrize("widget_cls", [FeaturesScatterWidget])
@@ -76,26 +75,35 @@ def assert_features_plot_changes(
         name: data + 1 for name, data in data[1]["features"].items()
     }
     viewer.add_points(data[0], **data[1])
-    assert_plot_changes(viewer, widget)
+    assert_plot_changes(viewer, widget, n_layers=1)
 
 
-def assert_plot_changes(viewer: Viewer, widget: NapariMPLWidget) -> None:
+def assert_plot_changes(
+    viewer: Viewer, widget: NapariMPLWidget, *, n_layers: int
+) -> None:
     """
     Assert that a widget plot changes when the layer selection
-    is changed. The passed viewer must already have two layers
+    is changed. The passed viewer must already have (2 * n_layers) layers
     loaded.
     """
-    # Select first layer
+    # Select first layer(s)
     viewer.layers.selection.clear()
-    viewer.layers.selection.add(viewer.layers[0])
+
+    for i in range(n_layers):
+        viewer.layers.selection.add(viewer.layers[i])
+    assert len(viewer.layers.selection) == n_layers
     fig1 = deepcopy(widget.figure)
 
-    # Re-selecting first layer should produce identical plot
+    # Re-selecting first layer(s) should produce identical plot
     viewer.layers.selection.clear()
-    viewer.layers.selection.add(viewer.layers[0])
+    for i in range(n_layers):
+        viewer.layers.selection.add(viewer.layers[i])
+    assert len(viewer.layers.selection) == n_layers
     assert_figures_equal(widget.figure, fig1)
 
-    # Plotting the second layer should produce a different plot
+    # Plotting the second layer(s) should produce a different plot
     viewer.layers.selection.clear()
-    viewer.layers.selection.add(viewer.layers[1])
+    for i in range(n_layers):
+        viewer.layers.selection.add(viewer.layers[n_layers + i])
+    assert len(viewer.layers.selection) == n_layers
     assert_figures_not_equal(widget.figure, fig1)