Add test for loader priorities

Author: aragilar

specutils/io/registers.py

@@ -126,11 +126,9 @@ def wrapper(*args, **kwargs):
 def custom_writer(label, dtype=Spectrum1D, priority=0):
     def decorator(func):
         if _astropy_has_priorities():
-            io_registry.register_writer(
-                label, Spectrum1D, func, priority=priority,
-            )
+            io_registry.register_writer(label, dtype, func, priority=priority)
         else:
-            io_registry.register_writer(label, Spectrum1D, func)
+            io_registry.register_writer(label, dtype, func)
 
         @wraps(func)
         def wrapper(*args, **kwargs):

specutils/tests/test_io.py

@@ -4,6 +4,7 @@
 This module tests SpecUtils io routines
 """
 
+from collections import Counter
 from specutils.io.parsing_utils import generic_spectrum_from_table # or something like that
 from astropy.io import registry
 from astropy.table import Table
@@ -15,7 +16,8 @@
 import warnings
 
 from specutils import Spectrum1D, SpectrumList
-from specutils.io import data_loader
+from specutils.io import data_loader, custom_writer
+from specutils.io.registers import _astropy_has_priorities
 
 
 def test_generic_spectrum_from_table(recwarn):
@@ -156,3 +158,46 @@ def reader(*args, **kwargs):
         # Clean up after ourselves
         registry.unregister_reader(format_name, datatype)
         registry.unregister_identifier(format_name, datatype)
+
+
+@pytest.mark.xfail(
+    not _astropy_has_priorities(),
+    reason="Test requires priorities to be implemented in astropy",
+    raises=registry.IORegistryError,
+)
+def test_loader_uses_priority(tmpdir):
+    counter = Counter()
+    fname = str(tmpdir.join('good.txt'))
+
+    with open(fname, 'w') as ff:
+        ff.write('\n')
+
+    def identifier(origin, *args, **kwargs):
+        fname = args[0]
+        return 'good' in fname
+
+    @data_loader("test_counting_loader1", identifier=identifier, priority=1)
+    def counting_loader1(*args, **kwargs):
+        counter["test1"] += 1
+        return Spectrum1D(
+            spectral_axis=np.arange(1,1.1,0.01)*u.AA,
+            flux=np.ones(len(wave))*1.e-14*u.Jy,
+        )
+
+    @data_loader("test_counting_loader2", identifier=identifier, priority=2)
+    def counting_loader2(*args, **kwargs):
+        counter["test2"] += 1
+        return Spectrum1D(
+            spectral_axis=np.arange(1,1.1,0.01)*u.AA,
+            flux=np.ones(len(wave))*1.e-14*u.Jy,
+        )
+
+    Spectrum1D.read(fname)
+    assert counter["test2"] == 1
+    assert counter["test1"] == 0
+
+    for datatype in [Spectrum1D, SpectrumList]:
+        registry.unregister_reader("test_counting_loader1", datatype)
+        registry.unregister_identifier("test_counting_loader1", datatype)
+        registry.unregister_reader("test_counting_loader2", datatype)
+        registry.unregister_identifier("test_counting_loader2", datatype)