From 1477ebed61c9f7811aae77829a5156e81ff8ea34 Mon Sep 17 00:00:00 2001
From: Derek Homeier <dhomeie@gwdg.de>
Date: Fri, 14 Oct 2022 20:00:12 +0200
Subject: [PATCH 1/5] TST: add test for exporting data with WCS

---
 .../image/qt/tests/test_python_export.py      | 30 +++++++++++++++----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/glue/viewers/image/qt/tests/test_python_export.py b/glue/viewers/image/qt/tests/test_python_export.py
index 3ee339619..14d928485 100644
--- a/glue/viewers/image/qt/tests/test_python_export.py
+++ b/glue/viewers/image/qt/tests/test_python_export.py
@@ -1,6 +1,8 @@
+import pytest
 import numpy as np
 import matplotlib.pyplot as plt
 from astropy.utils import NumpyRNGContext
+from astropy.wcs import WCS
 
 from glue.core import Data, DataCollection
 from glue.app.qt.application import GlueApplication
@@ -14,7 +16,9 @@ def setup_method(self, method):
 
         with NumpyRNGContext(12345):
             self.data = Data(cube=np.random.random((30, 50, 20)))
-        self.data_collection = DataCollection([self.data])
+        # Create data version with WCS coordinates
+        self.data_wcs = Data(label='cube', cube=self.data['cube'], coords=WCS(naxis=3))
+        self.data_collection = DataCollection([self.data, self.data_wcs])
         self.app = GlueApplication(self.data_collection)
         self.viewer = self.app.new_data_viewer(ImageViewer)
         self.viewer.add_data(self.data)
@@ -34,19 +38,35 @@ def teardown_method(self, method):
     def assert_same(self, tmpdir, tol=0.1):
         BaseTestExportPython.assert_same(self, tmpdir, tol=tol)
 
-    def test_simple(self, tmpdir):
+    @pytest.mark.parametrize('wcs', [False, True])
+    def test_simple(self, tmpdir, wcs):
+        if wcs:
+            self.viewer = self.app.new_data_viewer(ImageViewer)
+            self.viewer.add_data(self.data_wcs)
         self.assert_same(tmpdir)
 
-    def test_simple_legend(self, tmpdir):
+    @pytest.mark.parametrize('wcs', [False, True])
+    def test_simple_legend(self, tmpdir, wcs):
+        if wcs:
+            self.viewer = self.app.new_data_viewer(ImageViewer)
+            self.viewer.add_data(self.data_wcs)
         self.viewer.state.show_legend = True
         self.assert_same(tmpdir)
 
-    def test_simple_att(self, tmpdir):
+    @pytest.mark.parametrize('wcs', [False, True])
+    def test_simple_att(self, tmpdir, wcs):
+        if wcs:
+            self.viewer = self.app.new_data_viewer(ImageViewer)
+            self.viewer.add_data(self.data_wcs)
         self.viewer.state.x_att = self.data.pixel_component_ids[1]
         self.viewer.state.y_att = self.data.pixel_component_ids[0]
         self.assert_same(tmpdir)
 
-    def test_simple_visual(self, tmpdir):
+    @pytest.mark.parametrize('wcs', [False, True])
+    def test_simple_visual(self, tmpdir, wcs):
+        if wcs:
+            self.viewer = self.app.new_data_viewer(ImageViewer)
+            self.viewer.add_data(self.data_wcs)
         self.viewer.state.legend.visible = True
         self.viewer.state.layers[0].cmap = plt.cm.RdBu
         self.viewer.state.layers[0].v_min = 0.2

From 75a75b878b9b00b760756da56dc807f1a3490786 Mon Sep 17 00:00:00 2001
From: Derek Homeier <dhomeie@gwdg.de>
Date: Fri, 14 Oct 2022 20:09:46 +0200
Subject: [PATCH 2/5] Fix mpl export `_script_header` to use WCS in data.coords

---
 glue/viewers/image/viewer.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/glue/viewers/image/viewer.py b/glue/viewers/image/viewer.py
index 7fceee9c8..7b202f599 100644
--- a/glue/viewers/image/viewer.py
+++ b/glue/viewers/image/viewer.py
@@ -214,24 +214,26 @@ def _script_header(self):
 
         script = ""
         script += "fig, ax = init_mpl(wcs=True)\n"
-        script += "ax.set_aspect('{0}')\n".format(self.state.aspect)
+        script += f"ax.set_aspect('{self.state.aspect}')\n"
 
         script += '\ncomposite = CompositeArray()\n'
-        script += "image = imshow(ax, composite, origin='lower', interpolation='nearest', aspect='{0}')\n\n".format(self.state.aspect)
+        script += f"image = imshow(ax, composite, origin='lower', interpolation='nearest', aspect='{self.state.aspect}')\n\n"
 
         dindex = self.session.data_collection.index(self.state.reference_data)
 
-        script += "ref_data = data_collection[{0}]\n".format(dindex)
+        script += f"ref_data = data_collection[{dindex}]\n"
 
         ref_coords = self.state.reference_data.coords
 
-        if hasattr(ref_coords, 'wcs'):
-            script += "ax.reset_wcs(slices={0}, wcs=ref_data.coords.wcs)\n".format(self.state.wcsaxes_slice)
+        if isinstance(ref_coords, WCS):
+            script += f"ax.reset_wcs(slices={self.state.wcsaxes_slice}, wcs=ref_data.coords)\n"
+        elif hasattr(ref_coords, 'wcs'):
+            script += f"ax.reset_wcs(slices={self.state.wcsaxes_slice}, wcs=ref_data.coords.wcs)\n"
         elif hasattr(ref_coords, 'wcsaxes_dict'):
             raise NotImplementedError()
         else:
             imports.append('from glue.viewers.image.viewer import get_identity_wcs')
-            script += "ax.reset_wcs(slices={0}, wcs=get_identity_wcs(ref_data.ndim))\n".format(self.state.wcsaxes_slice)
+            script += f"ax.reset_wcs(slices={self.state.wcsaxes_slice}, wcs=get_identity_wcs(ref_data.ndim))\n"
 
         script += "# for the legend\n"
         script += "legend_handles = []\n"

From b480316d6222fcb026e9a80f54b0f142cd8101ca Mon Sep 17 00:00:00 2001
From: Derek Homeier <dhomeie@gwdg.de>
Date: Fri, 14 Oct 2022 23:34:28 +0200
Subject: [PATCH 3/5] Add more export tests with WCS

---
 .../image/qt/tests/test_python_export.py      | 38 ++++++++++++++-----
 1 file changed, 28 insertions(+), 10 deletions(-)

diff --git a/glue/viewers/image/qt/tests/test_python_export.py b/glue/viewers/image/qt/tests/test_python_export.py
index 14d928485..72270d19a 100644
--- a/glue/viewers/image/qt/tests/test_python_export.py
+++ b/glue/viewers/image/qt/tests/test_python_export.py
@@ -41,23 +41,23 @@ def assert_same(self, tmpdir, tol=0.1):
     @pytest.mark.parametrize('wcs', [False, True])
     def test_simple(self, tmpdir, wcs):
         if wcs:
-            self.viewer = self.app.new_data_viewer(ImageViewer)
             self.viewer.add_data(self.data_wcs)
+            self.viewer.remove_data(self.data)
         self.assert_same(tmpdir)
 
     @pytest.mark.parametrize('wcs', [False, True])
     def test_simple_legend(self, tmpdir, wcs):
         if wcs:
-            self.viewer = self.app.new_data_viewer(ImageViewer)
             self.viewer.add_data(self.data_wcs)
+            self.viewer.remove_data(self.data)
         self.viewer.state.show_legend = True
         self.assert_same(tmpdir)
 
     @pytest.mark.parametrize('wcs', [False, True])
     def test_simple_att(self, tmpdir, wcs):
         if wcs:
-            self.viewer = self.app.new_data_viewer(ImageViewer)
             self.viewer.add_data(self.data_wcs)
+            self.viewer.remove_data(self.data)
         self.viewer.state.x_att = self.data.pixel_component_ids[1]
         self.viewer.state.y_att = self.data.pixel_component_ids[0]
         self.assert_same(tmpdir)
@@ -65,8 +65,8 @@ def test_simple_att(self, tmpdir, wcs):
     @pytest.mark.parametrize('wcs', [False, True])
     def test_simple_visual(self, tmpdir, wcs):
         if wcs:
-            self.viewer = self.app.new_data_viewer(ImageViewer)
             self.viewer.add_data(self.data_wcs)
+            self.viewer.remove_data(self.data)
         self.viewer.state.legend.visible = True
         self.viewer.state.layers[0].cmap = plt.cm.RdBu
         self.viewer.state.layers[0].v_min = 0.2
@@ -91,17 +91,35 @@ def test_subset(self, tmpdir):
         self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
         self.assert_same(tmpdir)
 
-    def test_subset_legend(self, tmpdir):
+    @pytest.mark.parametrize('wcs', [False, True])
+    def test_subset_legend(self, tmpdir, wcs):
+        if wcs:
+            self.viewer.add_data(self.data_wcs)
+            self.viewer.remove_data(self.data)
+            self.data_collection.new_subset_group('mysubset', self.data_wcs.id['cube'] > 0.5)
+        else:
+            self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
         self.viewer.state.legend.visible = True
-        self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
         self.assert_same(tmpdir, tol=0.15)  # transparency and such
 
-    def test_subset_slice(self, tmpdir):
-        self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
+    @pytest.mark.parametrize('wcs', [False, True])
+    def test_subset_slice(self, tmpdir, wcs):
+        if wcs:
+            self.viewer.add_data(self.data_wcs)
+            self.viewer.remove_data(self.data)
+            self.data_collection.new_subset_group('mysubset', self.data_wcs.id['cube'] > 0.5)
+        else:
+            self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
         self.test_slice(tmpdir)
 
-    def test_subset_transposed(self, tmpdir):
+    @pytest.mark.parametrize('wcs', [False, True])
+    def test_subset_transposed(self, tmpdir, wcs):
+        if wcs:
+            self.viewer.add_data(self.data_wcs)
+            self.viewer.remove_data(self.data)
+            self.data_collection.new_subset_group('mysubset', self.data_wcs.id['cube'] > 0.5)
+        else:
+            self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
         self.viewer.state.x_att = self.data.pixel_component_ids[0]
         self.viewer.state.y_att = self.data.pixel_component_ids[1]
-        self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
         self.assert_same(tmpdir)

From a7dae6ef8c5f9ce9fdf620518a17f07003a27a86 Mon Sep 17 00:00:00 2001
From: Derek Homeier <dhomeie@gwdg.de>
Date: Sat, 15 Oct 2022 16:56:15 +0200
Subject: [PATCH 4/5] Simplify coords treatment; add tests for
 AffineCoordinates

---
 .../image/qt/tests/test_python_export.py      | 125 +++++++++++-------
 glue/viewers/image/viewer.py                  |  15 +--
 2 files changed, 83 insertions(+), 57 deletions(-)

diff --git a/glue/viewers/image/qt/tests/test_python_export.py b/glue/viewers/image/qt/tests/test_python_export.py
index 72270d19a..64b24e487 100644
--- a/glue/viewers/image/qt/tests/test_python_export.py
+++ b/glue/viewers/image/qt/tests/test_python_export.py
@@ -5,6 +5,7 @@
 from astropy.wcs import WCS
 
 from glue.core import Data, DataCollection
+from glue.core.coordinates import AffineCoordinates
 from glue.app.qt.application import GlueApplication
 from glue.viewers.image.qt import ImageViewer
 from glue.viewers.matplotlib.qt.tests.test_python_export import BaseTestExportPython
@@ -16,9 +17,13 @@ def setup_method(self, method):
 
         with NumpyRNGContext(12345):
             self.data = Data(cube=np.random.random((30, 50, 20)))
-        # Create data version with WCS coordinates
+        # Create data versions with WCS and affine coordinates
+        matrix = np.array([[2, 3, 4, -1], [1, 2, 2, 2], [1, 1, 1, -2], [0, 0, 0, 1]])
+        affine = AffineCoordinates(matrix, units=['Mm', 'Mm', 'km'], labels=['xw', 'yw', 'zw'])
+
         self.data_wcs = Data(label='cube', cube=self.data['cube'], coords=WCS(naxis=3))
-        self.data_collection = DataCollection([self.data, self.data_wcs])
+        self.data_affine = Data(label='cube', cube=self.data['cube'], coords=affine)
+        self.data_collection = DataCollection([self.data, self.data_wcs, self.data_affine])
         self.app = GlueApplication(self.data_collection)
         self.viewer = self.app.new_data_viewer(ImageViewer)
         self.viewer.add_data(self.data)
@@ -38,34 +43,39 @@ def teardown_method(self, method):
     def assert_same(self, tmpdir, tol=0.1):
         BaseTestExportPython.assert_same(self, tmpdir, tol=tol)
 
-    @pytest.mark.parametrize('wcs', [False, True])
-    def test_simple(self, tmpdir, wcs):
-        if wcs:
-            self.viewer.add_data(self.data_wcs)
+    @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
+    def test_simple(self, tmpdir, coords):
+        if coords is not None:
+            self.viewer.add_data(getattr(self, f'data_{coords}'))
             self.viewer.remove_data(self.data)
         self.assert_same(tmpdir)
 
-    @pytest.mark.parametrize('wcs', [False, True])
-    def test_simple_legend(self, tmpdir, wcs):
-        if wcs:
-            self.viewer.add_data(self.data_wcs)
+    @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
+    def test_simple_legend(self, tmpdir, coords):
+        if coords is not None:
+            self.viewer.add_data(getattr(self, f'data_{coords}'))
             self.viewer.remove_data(self.data)
         self.viewer.state.show_legend = True
         self.assert_same(tmpdir)
 
-    @pytest.mark.parametrize('wcs', [False, True])
-    def test_simple_att(self, tmpdir, wcs):
-        if wcs:
-            self.viewer.add_data(self.data_wcs)
+    @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
+    def test_simple_att(self, tmpdir, coords):
+        if coords is None:
+            data = self.data
+        else:
+            data = getattr(self, f'data_{coords}')
+            self.viewer.add_data(data)
             self.viewer.remove_data(self.data)
-        self.viewer.state.x_att = self.data.pixel_component_ids[1]
-        self.viewer.state.y_att = self.data.pixel_component_ids[0]
+        self.viewer.state.x_att = data.pixel_component_ids[1]
+        self.viewer.state.y_att = data.pixel_component_ids[0]
+        if coords == 'affine':
+            pytest.xfail('Known issue with axis label rendering')
         self.assert_same(tmpdir)
 
-    @pytest.mark.parametrize('wcs', [False, True])
-    def test_simple_visual(self, tmpdir, wcs):
-        if wcs:
-            self.viewer.add_data(self.data_wcs)
+    @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
+    def test_simple_visual(self, tmpdir, coords):
+        if coords is not None:
+            self.viewer.add_data(getattr(self, f'data_{coords}'))
             self.viewer.remove_data(self.data)
         self.viewer.state.legend.visible = True
         self.viewer.state.layers[0].cmap = plt.cm.RdBu
@@ -77,49 +87,70 @@ def test_simple_visual(self, tmpdir, wcs):
         self.viewer.state.layers[0].bias = 0.6
         self.assert_same(tmpdir)
 
-    def test_slice(self, tmpdir):
-        self.viewer.state.x_att = self.data.pixel_component_ids[1]
-        self.viewer.state.y_att = self.data.pixel_component_ids[0]
+    @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
+    def test_slice(self, tmpdir, coords):
+        if coords is None:
+            data = self.data
+        else:
+            data = getattr(self, f'data_{coords}')
+            self.viewer.add_data(data)
+            self.viewer.remove_data(self.data)
+
+        self.viewer.state.x_att = data.pixel_component_ids[1]
+        self.viewer.state.y_att = data.pixel_component_ids[0]
         self.viewer.state.slices = (2, 3, 4)
+        if coords == 'affine':
+            pytest.xfail('Known issue with axis label rendering')
         self.assert_same(tmpdir)
 
-    def test_aspect(self, tmpdir):
+    @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
+    def test_aspect(self, tmpdir, coords):
+        if coords is not None:
+            self.viewer.add_data(getattr(self, f'data_{coords}'))
+            self.viewer.remove_data(self.data)
         self.viewer.state.aspect = 'auto'
         self.assert_same(tmpdir)
 
-    def test_subset(self, tmpdir):
+    @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
+    def test_subset(self, tmpdir, coords):
+        if coords is not None:
+            self.viewer.add_data(getattr(self, f'data_{coords}'))
+            self.viewer.remove_data(self.data)
         self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
         self.assert_same(tmpdir)
 
-    @pytest.mark.parametrize('wcs', [False, True])
-    def test_subset_legend(self, tmpdir, wcs):
-        if wcs:
-            self.viewer.add_data(self.data_wcs)
-            self.viewer.remove_data(self.data)
-            self.data_collection.new_subset_group('mysubset', self.data_wcs.id['cube'] > 0.5)
+    @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
+    def test_subset_legend(self, tmpdir, coords):
+        if coords is None:
+            data = self.data
         else:
-            self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
+            data = getattr(self, f'data_{coords}')
+            self.viewer.add_data(data)
+            self.viewer.remove_data(self.data)
+        self.data_collection.new_subset_group('mysubset', data.id['cube'] > 0.5)
         self.viewer.state.legend.visible = True
         self.assert_same(tmpdir, tol=0.15)  # transparency and such
 
-    @pytest.mark.parametrize('wcs', [False, True])
-    def test_subset_slice(self, tmpdir, wcs):
-        if wcs:
-            self.viewer.add_data(self.data_wcs)
-            self.viewer.remove_data(self.data)
-            self.data_collection.new_subset_group('mysubset', self.data_wcs.id['cube'] > 0.5)
+    @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
+    def test_subset_slice(self, tmpdir, coords):
+        if coords is None:
+            data = self.data
         else:
-            self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
-        self.test_slice(tmpdir)
-
-    @pytest.mark.parametrize('wcs', [False, True])
-    def test_subset_transposed(self, tmpdir, wcs):
-        if wcs:
-            self.viewer.add_data(self.data_wcs)
+            data = getattr(self, f'data_{coords}')
+            self.viewer.add_data(data)
             self.viewer.remove_data(self.data)
-            self.data_collection.new_subset_group('mysubset', self.data_wcs.id['cube'] > 0.5)
+        self.data_collection.new_subset_group('mysubset', data.id['cube'] > 0.5)
+        self.test_slice(tmpdir, coords)
+
+    @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
+    def test_subset_transposed(self, tmpdir, coords):
+        if coords is None:
+            data = self.data
         else:
-            self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
+            data = getattr(self, f'data_{coords}')
+            self.viewer.add_data(data)
+            self.viewer.remove_data(self.data)
+        self.data_collection.new_subset_group('mysubset', data.id['cube'] > 0.5)
         self.viewer.state.x_att = self.data.pixel_component_ids[0]
         self.viewer.state.y_att = self.data.pixel_component_ids[1]
         self.assert_same(tmpdir)
diff --git a/glue/viewers/image/viewer.py b/glue/viewers/image/viewer.py
index 7b202f599..c4ffb3b0f 100644
--- a/glue/viewers/image/viewer.py
+++ b/glue/viewers/image/viewer.py
@@ -223,17 +223,12 @@ def _script_header(self):
 
         script += f"ref_data = data_collection[{dindex}]\n"
 
-        ref_coords = self.state.reference_data.coords
-
-        if isinstance(ref_coords, WCS):
-            script += f"ax.reset_wcs(slices={self.state.wcsaxes_slice}, wcs=ref_data.coords)\n"
-        elif hasattr(ref_coords, 'wcs'):
-            script += f"ax.reset_wcs(slices={self.state.wcsaxes_slice}, wcs=ref_data.coords.wcs)\n"
-        elif hasattr(ref_coords, 'wcsaxes_dict'):
-            raise NotImplementedError()
-        else:
+        if isinstance(self.state.reference_data.coords, (LegacyCoordinates, type(None))):
             imports.append('from glue.viewers.image.viewer import get_identity_wcs')
-            script += f"ax.reset_wcs(slices={self.state.wcsaxes_slice}, wcs=get_identity_wcs(ref_data.ndim))\n"
+            ref_wcs = "get_identity_wcs(ref_data.ndim)"
+        else:
+            ref_wcs = "ref_data.coords"
+        script += f"ax.reset_wcs(slices={self.state.wcsaxes_slice}, wcs={ref_wcs})\n"
 
         script += "# for the legend\n"
         script += "legend_handles = []\n"

From 0cbfbd1e5706828cd60d1fbd1ad602b9709f515d Mon Sep 17 00:00:00 2001
From: Derek Homeier <dhomeie@gwdg.de>
Date: Thu, 27 Oct 2022 15:49:28 +0200
Subject: [PATCH 5/5] TST: use helper function to load alternative
 `Coordinates` datasets

---
 .../image/qt/tests/test_python_export.py      | 74 ++++++-------------
 1 file changed, 22 insertions(+), 52 deletions(-)

diff --git a/glue/viewers/image/qt/tests/test_python_export.py b/glue/viewers/image/qt/tests/test_python_export.py
index 64b24e487..cfee1c89a 100644
--- a/glue/viewers/image/qt/tests/test_python_export.py
+++ b/glue/viewers/image/qt/tests/test_python_export.py
@@ -43,40 +43,34 @@ def teardown_method(self, method):
     def assert_same(self, tmpdir, tol=0.1):
         BaseTestExportPython.assert_same(self, tmpdir, tol=tol)
 
-    @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
-    def test_simple(self, tmpdir, coords):
+    def viewer_load(self, coords):
         if coords is not None:
             self.viewer.add_data(getattr(self, f'data_{coords}'))
             self.viewer.remove_data(self.data)
+
+    @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
+    def test_simple(self, tmpdir, coords):
+        self.viewer_load(coords)
         self.assert_same(tmpdir)
 
     @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
     def test_simple_legend(self, tmpdir, coords):
-        if coords is not None:
-            self.viewer.add_data(getattr(self, f'data_{coords}'))
-            self.viewer.remove_data(self.data)
+        self.viewer_load(coords)
         self.viewer.state.show_legend = True
         self.assert_same(tmpdir)
 
     @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
     def test_simple_att(self, tmpdir, coords):
-        if coords is None:
-            data = self.data
-        else:
-            data = getattr(self, f'data_{coords}')
-            self.viewer.add_data(data)
-            self.viewer.remove_data(self.data)
-        self.viewer.state.x_att = data.pixel_component_ids[1]
-        self.viewer.state.y_att = data.pixel_component_ids[0]
+        self.viewer_load(coords)
+        self.viewer.state.x_att = self.viewer.state.reference_data.pixel_component_ids[1]
+        self.viewer.state.y_att = self.viewer.state.reference_data.pixel_component_ids[0]
         if coords == 'affine':
             pytest.xfail('Known issue with axis label rendering')
         self.assert_same(tmpdir)
 
     @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
     def test_simple_visual(self, tmpdir, coords):
-        if coords is not None:
-            self.viewer.add_data(getattr(self, f'data_{coords}'))
-            self.viewer.remove_data(self.data)
+        self.viewer_load(coords)
         self.viewer.state.legend.visible = True
         self.viewer.state.layers[0].cmap = plt.cm.RdBu
         self.viewer.state.layers[0].v_min = 0.2
@@ -89,15 +83,9 @@ def test_simple_visual(self, tmpdir, coords):
 
     @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
     def test_slice(self, tmpdir, coords):
-        if coords is None:
-            data = self.data
-        else:
-            data = getattr(self, f'data_{coords}')
-            self.viewer.add_data(data)
-            self.viewer.remove_data(self.data)
-
-        self.viewer.state.x_att = data.pixel_component_ids[1]
-        self.viewer.state.y_att = data.pixel_component_ids[0]
+        self.viewer_load(coords)
+        self.viewer.state.x_att = self.viewer.state.reference_data.pixel_component_ids[1]
+        self.viewer.state.y_att = self.viewer.state.reference_data.pixel_component_ids[0]
         self.viewer.state.slices = (2, 3, 4)
         if coords == 'affine':
             pytest.xfail('Known issue with axis label rendering')
@@ -105,52 +93,34 @@ def test_slice(self, tmpdir, coords):
 
     @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
     def test_aspect(self, tmpdir, coords):
-        if coords is not None:
-            self.viewer.add_data(getattr(self, f'data_{coords}'))
-            self.viewer.remove_data(self.data)
+        self.viewer_load(coords)
         self.viewer.state.aspect = 'auto'
         self.assert_same(tmpdir)
 
     @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
     def test_subset(self, tmpdir, coords):
-        if coords is not None:
-            self.viewer.add_data(getattr(self, f'data_{coords}'))
-            self.viewer.remove_data(self.data)
+        self.viewer_load(coords)
         self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
         self.assert_same(tmpdir)
 
     @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
     def test_subset_legend(self, tmpdir, coords):
-        if coords is None:
-            data = self.data
-        else:
-            data = getattr(self, f'data_{coords}')
-            self.viewer.add_data(data)
-            self.viewer.remove_data(self.data)
-        self.data_collection.new_subset_group('mysubset', data.id['cube'] > 0.5)
+        self.viewer_load(coords)
+        self.data_collection.new_subset_group('mysubset',
+                                              self.viewer.state.reference_data.id['cube'] > 0.5)
         self.viewer.state.legend.visible = True
         self.assert_same(tmpdir, tol=0.15)  # transparency and such
 
     @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
     def test_subset_slice(self, tmpdir, coords):
-        if coords is None:
-            data = self.data
-        else:
-            data = getattr(self, f'data_{coords}')
-            self.viewer.add_data(data)
-            self.viewer.remove_data(self.data)
-        self.data_collection.new_subset_group('mysubset', data.id['cube'] > 0.5)
+        self.viewer_load(coords)
+        self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
         self.test_slice(tmpdir, coords)
 
     @pytest.mark.parametrize('coords', [None, 'wcs', 'affine'])
     def test_subset_transposed(self, tmpdir, coords):
-        if coords is None:
-            data = self.data
-        else:
-            data = getattr(self, f'data_{coords}')
-            self.viewer.add_data(data)
-            self.viewer.remove_data(self.data)
-        self.data_collection.new_subset_group('mysubset', data.id['cube'] > 0.5)
+        self.viewer_load(coords)
+        self.data_collection.new_subset_group('mysubset', self.data.id['cube'] > 0.5)
         self.viewer.state.x_att = self.data.pixel_component_ids[0]
         self.viewer.state.y_att = self.data.pixel_component_ids[1]
         self.assert_same(tmpdir)