TST: put collage tests under a class

tacaswell · tacaswell · commit 825197a6ce07 · 2020-03-01T23:37:45.000-05:00
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
@@ -501,64 +501,65 @@ def test_removed_axis():
     axs[0].remove()
     fig.canvas.draw()
 
-@check_figures_equal()
-@pytest.mark.parametrize('x', [
-    [["A", "A", "B"], ["C", "D", "B"]],
-    [[1, 1, 2], [3, 4, 2]]]
-)
-def test_collage_basic(fig_test, fig_ref, x):
-    grid_axes = fig_test.build_grid(x)
 
-    for k, ax in grid_axes.items():
-        ax.set_title(k)
+class TestCollage:
+    @check_figures_equal()
+    @pytest.mark.parametrize('x', [
+        [["A", "A", "B"], ["C", "D", "B"]],
+        [[1, 1, 2], [3, 4, 2]]]
+    )
+    def test_basic(self, fig_test, fig_ref, x):
+        grid_axes = fig_test.build_grid(x)
 
-    labels = sorted(np.unique(x))
+        for k, ax in grid_axes.items():
+            ax.set_title(k)
 
-    assert len(labels) == len(grid_axes)
+        labels = sorted(np.unique(x))
 
-    gs = fig_ref.add_gridspec(2, 3)
-    axA = fig_ref.add_subplot(gs[:1, :2])
-    axA.set_title(labels[0])
+        assert len(labels) == len(grid_axes)
 
-    axB = fig_ref.add_subplot(gs[:, 2])
-    axB.set_title(labels[1])
+        gs = fig_ref.add_gridspec(2, 3)
+        axA = fig_ref.add_subplot(gs[:1, :2])
+        axA.set_title(labels[0])
 
-    axC = fig_ref.add_subplot(gs[1, 0])
-    axC.set_title(labels[2])
+        axB = fig_ref.add_subplot(gs[:, 2])
+        axB.set_title(labels[1])
 
-    axD = fig_ref.add_subplot(gs[1, 1])
-    axD.set_title(labels[3])
+        axC = fig_ref.add_subplot(gs[1, 0])
+        axC.set_title(labels[2])
 
+        axD = fig_ref.add_subplot(gs[1, 1])
+        axD.set_title(labels[3])
 
-@check_figures_equal(tol=0.005)
-def test_collage_nested(fig_test, fig_ref):
-    x = [["A", "B", "B"], ["C", "C", "D"]]
+    @check_figures_equal(tol=0.005)
+    def test_nested(self, fig_test, fig_ref):
+        x = [["A", "B", "B"], ["C", "C", "D"]]
 
-    y = [["F"], [x]]
+        y = [["F"], [x]]
 
-    grid_axes = fig_test.build_grid(y)
+        grid_axes = fig_test.build_grid(y)
 
-    for k, ax in grid_axes.items():
-        ax.set_title(k)
+        for k, ax in grid_axes.items():
+            ax.set_title(k)
 
-    gs = fig_ref.add_gridspec(2, 1)
+        gs = fig_ref.add_gridspec(2, 1)
 
-    gs_n = gs[1, 0].subgridspec(2, 3)
+        gs_n = gs[1, 0].subgridspec(2, 3)
 
-    axA = fig_ref.add_subplot(gs_n[0, 0])
-    axA.set_title("A")
+        axA = fig_ref.add_subplot(gs_n[0, 0])
+        axA.set_title("A")
 
-    axB = fig_ref.add_subplot(gs_n[0, 1:])
-    axB.set_title("B")
+        axB = fig_ref.add_subplot(gs_n[0, 1:])
+        axB.set_title("B")
 
-    axC = fig_ref.add_subplot(gs_n[1, :2])
-    axC.set_title("C")
+        axC = fig_ref.add_subplot(gs_n[1, :2])
+        axC.set_title("C")
 
-    axD = fig_ref.add_subplot(gs_n[1, 2])
-    axD.set_title("D")
+        axD = fig_ref.add_subplot(gs_n[1, 2])
+        axD.set_title("D")
 
-    axF = fig_ref.add_subplot(gs[0, 0])
-    axF.set_title("F")
+        axF = fig_ref.add_subplot(gs[0, 0])
+        axF.set_title("F")
 
-    fig_test.tight_layout()
-    fig_ref.tight_layout()
+        fig_test.tight_layout()
+        fig_ref.tight_layout()
