Add tests

4 files changed, 94 insertions, 17 deletions

diff --git a/stwcs/tests/test_altwcs.py b/stwcs/tests/test_altwcs.py
index d646106..fa7eb9b 100644
--- a/stwcs/tests/test_altwcs.py
+++ b/stwcs/tests/test_altwcs.py
@@ -117,7 +117,6 @@ class TestAltWCS(object):
 
     def test_restore_wcs_from_to(self):
         # test restore from ... to ...
-        #altwcs.archiveWCS(self.acs_file, ext=[('SCI',1), ('SCI',2)], wcskey='T')
         pyfits.setval(self.acs_file, ext=('SCI', 1), keyword='CRVAL1', value=1)
         pyfits.setval(self.acs_file, ext=('SCI', 2), keyword='CRVAL1', value=1)
         f = pyfits.open(self.acs_file, mode='update')
@@ -140,9 +139,7 @@ class TestAltWCS(object):
         compare_wcs(w3, w1o, exclude_keywords=['ctype'])
 
     def test_delete_wcs(self):
-        #altwcs.archiveWCS(self.acs_file, ext=1, wcskey='Z')
         altwcs.deleteWCS(self.acs_file, ext=1, wcskey='Z')
-        #utils.assert_raises(KeyError, wcsutil.HSTWCS, self.acs_file, ext=1, wcskey='Z')
         with pytest.raises(KeyError):
             wcsutil.HSTWCS(self.acs_file, ext=1, wcskey='Z')
 
diff --git a/stwcs/tests/test_headerlet.py b/stwcs/tests/test_headerlet.py
index f4fa050..462b7e4 100644
--- a/stwcs/tests/test_headerlet.py
+++ b/stwcs/tests/test_headerlet.py
@@ -8,7 +8,7 @@ from ..wcsutil import headerlet, wcsdiff
 from ..wcsutil import HSTWCS
 import numpy as np
 from numpy.testing import utils
-from nose.tools import *
+import pytest
 
 from . import data
 data_path = os.path.split(os.path.abspath(data.__file__))[0]
@@ -100,7 +100,6 @@ class TestCreateHeaderlet(object):
         assert(wcsdiff.is_wcs_identical(self.simple_file, self.headerlet_name,
                                         [0], [1], verbose=True)[0])
 
-    @raises(KeyError)
     def test_no_HDRNAME_no_WCSNAME(self):
         """
         Test create_headerlet stepping through all
@@ -110,7 +109,8 @@ class TestCreateHeaderlet(object):
         shutil.copyfile(self.comp_file, newf)
         fits.delval(newf, 'HDRNAME', ext=1)
         fits.delval(newf, 'WCSNAME', ext=1)
-        hlet = headerlet.create_headerlet(newf)
+        with pytest.raises(KeyError):
+            hlet = headerlet.create_headerlet(newf)
 
     def test1SciExt(self):
         """
@@ -176,23 +176,24 @@ class TestApplyHeaderlet:
         hlet.apply_as_primary('comp.fits')
     """
 
-    @raises(ValueError)
+
     def testWrongSIPModel(self):
         hlet = headerlet.create_headerlet(self.comp_file, hdrname='test1',
                                           sipname='WRONG')
-        hlet.apply_as_primary(self.comp_file)
+        with pytest.raises(ValueError):
+            hlet.apply_as_primary(self.comp_file)
 
-    @raises(ValueError)
     def testWrongNPOLModel(self):
         hlet = headerlet.create_headerlet(self.comp_file, hdrname='test1',
                                           npolfile='WRONG')
-        hlet.apply_as_primary(self.comp_file)
+        with pytest.raises(ValueError):
+            hlet.apply_as_primary(self.comp_file)
 
-    @raises(ValueError)
     def testWrongD2IMModel(self):
         hlet = headerlet.create_headerlet(self.comp_file, hdrname='test1',
                                           d2imfile='WRONG')
-        hlet.apply_as_primary(self.comp_file)
+        with pytest.raises(ValueError):
+            hlet.apply_as_primary(self.comp_file)
 
     def test_apply_as_primary_method(self):
         hlet = headerlet.create_headerlet(self.comp_file, hdrname='test2')
diff --git a/stwcs/tests/test_updatewcs.py b/stwcs/tests/test_updatewcs.py
index b7e75c8..4afd7cd 100644
--- a/stwcs/tests/test_updatewcs.py
+++ b/stwcs/tests/test_updatewcs.py
@@ -296,12 +296,69 @@ def test_add_radesys():
     fits.setval(acs_file, ext=0, keyword="NPOLFILE", value=npol_file)
     fits.setval(acs_file, ext=0, keyword="D2IMFILE", value=d2imfile)
 
-    #shutil.copyfile('orig/ibof01ahq_flt.fits', './ibof01ahq_flt.fits')
     updatewcs.updatewcs(acs_file)
-    # updatewcs.updatewcs('ibof01ahq_flt.fits')
     for ext in [('SCI', 1), ('SCI', 2)]:
         hdr = fits.getheader(acs_file, ext)
         assert hdr['RADESYS'] == 'FK5'
 
-    #hdr = fits.getheader('ibof01ahq_flt.fits', ext=('SCI', 1))
-    #assert hdr['RADESYS'] == 'ICRS'
+def test_update_d2im_distortion():
+    acs_orig_file = get_filepath('j94f05bgq_flt.fits')
+    current_dir = os.path.abspath(os.path.curdir)
+    acs_file = get_filepath('j94f05bgq_flt.fits', current_dir)
+
+    idctab = get_filepath('postsm4_idc.fits')
+    npol_file = get_filepath('qbu16424j_npl.fits')
+    d2imfile = get_filepath('new_wfc_d2i.fits')
+    newd2im = get_filepath('new_wfc_d2i.fits', current_dir)
+    print('newd2im', newd2im)
+    try:
+        os.remove(acs_file)
+    except OSError:
+        pass
+    shutil.copyfile(acs_orig_file, acs_file)
+    fits.setval(acs_file, ext=0, keyword="IDCTAB", value=idctab)
+    fits.setval(acs_file, ext=0, keyword="NPOLFILE", value=npol_file)
+    fits.setval(acs_file, ext=0, keyword="D2IMFILE", value=d2imfile)
+    updatewcs.updatewcs(acs_file)
+    d2imerr1 = fits.getval(acs_file, ext=1, keyword='D2IMERR1')
+    d2imerr4 = fits.getval(acs_file, ext=4, keyword='D2IMERR1')
+    shutil.copyfile(d2imfile, newd2im)
+    newf = fits.open(newd2im, mode='update')
+    for ext in newf[1:]:
+        ext.data *= 100
+    newf.close()
+    fits.setval(acs_file, keyword="D2IMFILE", value=newd2im)
+    updatewcs.updatewcs(acs_file)
+    nd2imerr1 = fits.getval(acs_file, ext=1, keyword='D2IMERR1')
+    nd2imerr4 = fits.getval(acs_file, ext=4, keyword='D2IMERR1')
+    assert np.isclose(d2imerr1 * 100, nd2imerr1)
+    assert np.isclose(d2imerr4 * 100, nd2imerr4)
+
+
+def test_apply_d2im():
+    from stwcs.updatewcs import apply_corrections as appc
+    acs_orig_file = get_filepath('j94f05bgq_flt.fits')
+    current_dir = os.path.abspath(os.path.curdir)
+    fname = get_filepath('j94f05bgq_flt.fits', current_dir)
+    d2imfile = get_filepath('new_wfc_d2i.fits')
+    try:
+        os.remove(fname)
+    except OSError:
+        pass
+    shutil.copyfile(acs_orig_file, fname)
+    fits.setval(fname, ext=0, keyword="D2IMFILE", value=d2imfile)
+    fits.setval(fname, ext=0, keyword="IDCTAB", value='N/A')
+    fits.setval(fname, ext=0, keyword="NPOLFILE", value='N/A')
+    # If D2IMEXT does not exist, the correction should be applied
+    assert appc.apply_d2im_correction(fname, d2imcorr=True)
+    updatewcs.updatewcs(fname)
+
+    # Test the case when D2IMFILE == D2IMEXT
+    assert not appc.apply_d2im_correction(fname, d2imcorr=True)
+    assert not appc.apply_d2im_correction(fname, d2imcorr=False)
+
+    fits.setval(fname, ext=0, keyword='D2IMFILE', value="N/A")
+    assert not appc.apply_d2im_correction(fname, d2imcorr=True)
+    # No D2IMFILE keyword in primary header
+    fits.delval(fname, ext=0, keyword='D2IMFILE')
+    assert not appc.apply_d2im_correction(fname, d2imcorr=True)
diff --git a/stwcs/updatewcs/apply_corrections.py b/stwcs/updatewcs/apply_corrections.py
index c554564..247dba0 100644
--- a/stwcs/updatewcs/apply_corrections.py
+++ b/stwcs/updatewcs/apply_corrections.py
@@ -212,6 +212,28 @@ def isOldStyleDGEO(fname, dgname):
 
 
 def apply_d2im_correction(fname, d2imcorr):
+    """
+    Logic to decide whether to apply the D2IM correction.
+
+    Parameters
+    ----------
+    fname : str
+        Science file name.
+    d2imcorr : bool
+        Flag indicating if D2IM is should be enabled if allowed.
+
+    Return
+    ------
+    applyD2IMCorr : bool
+        Flag whether to apply the correction.
+
+    The D2IM correction is applied to a science file if it is in the
+    allowed corrections for the instrument. The name of the file
+    with the correction is saved in the ``D2IMFILE`` keyword in the
+    primary header. When the correction is applied the name of the
+    file is saved in the ``D2IMEXT`` keyword in the 1st extension header.
+
+    """
     applyD2IMCorr = True
     if not d2imcorr:
         logger.info("D2IM correction not requested - not applying it.")
@@ -221,7 +243,7 @@ def apply_d2im_correction(fname, d2imcorr):
         fd2im0 = fits.getval(fname, 'D2IMFILE')
     except KeyError:
         logger.info("D2IMFILE keyword is missing - D2IM correction will not be applied.")
-
+        return False
     if fd2im0 == 'N/A':
         utils.remove_distortion(fname, "D2IMFILE")
         return False