diff --git a/specreduce/tests/test_align_along_trace.py b/specreduce/tests/test_align_along_trace.py
index fb6a9ad..c96ff97 100644
--- a/specreduce/tests/test_align_along_trace.py
+++ b/specreduce/tests/test_align_along_trace.py
@@ -3,7 +3,7 @@
 import astropy.units as u
 
 from specreduce.tracing import ArrayTrace
-from specreduce.utils import align_spectrum_along_trace
+from specreduce.utils import align_2d_spectrum_along_trace
 
 
 def mk_test_image():
@@ -18,19 +18,19 @@ def mk_test_image():
 def test_align_spectrum_along_trace_bad_input():
     image, trace = mk_test_image()
     with pytest.raises(ValueError, match='Unre'):
-        im = align_spectrum_along_trace(image, None)   # noqa
+        im = align_2d_spectrum_along_trace(image, None)   # noqa
 
     with pytest.raises(ValueError, match='method must be'):
-        im = align_spectrum_along_trace(image, trace, method='int')   # noqa
+        im = align_2d_spectrum_along_trace(image, trace, method='int')   # noqa
 
     with pytest.raises(ValueError, match='Spectral axis length'):
-        im = align_spectrum_along_trace(image.T, trace, method='interpolate', disp_axis=0)   # noqa
+        im = align_2d_spectrum_along_trace(image.T, trace, method='interpolate', disp_axis=0)   # noqa
 
     with pytest.raises(ValueError, match='Displacement axis must be'):
-        im = align_spectrum_along_trace(image, trace, disp_axis=2)  # noqa
+        im = align_2d_spectrum_along_trace(image, trace, disp_axis=2)  # noqa
 
     with pytest.raises(ValueError, match='The number of image dimensions must be'):
-        im = align_spectrum_along_trace(np.zeros((3, 6, 9)), trace)  # noqa
+        im = align_2d_spectrum_along_trace(np.zeros((3, 6, 9)), trace)  # noqa
 
 
 @pytest.mark.parametrize("method, truth_data, truth_mask, truth_ucty",
@@ -50,7 +50,7 @@ def test_align_spectrum_along_trace_bad_input():
                          ids=('method=interpolate', 'method=shift'))
 def test_align_spectrum_along_trace(method, truth_data, truth_mask, truth_ucty):
     image, trace = mk_test_image()
-    im = align_spectrum_along_trace(image, trace, method=method)
+    im = align_2d_spectrum_along_trace(image, trace, method=method)
     assert im.shape == image.shape
     assert im.unit == u.DN
     assert im.uncertainty.uncertainty_type == 'var'