Unit test added for wcs_functions.py

drizzlepac/tests/test_wcs_functions.py

@@ -0,0 +1,23 @@
+import os
+import pytest
+import numpy as np
+from astropy.wcs import WCS
+from astropy.io import fits
+from drizzlepac.wcs_functions import make_mosaic_wcs
+
+abs_path = 'drizzlepac/tests/'
+
+# checks if function works with single filenames and lists of filenames
+# does not check if wcs is different, but this was confirmed independently. 
+@pytest.mark.parametrize(
+    "filename_array",
+    [
+        (abs_path+"sample_drc_header.fits"),
+        (abs_path+"sample_flc_header.fits"),
+        ([abs_path+"sample_flc_header.fits", abs_path+"sample_flc_header.fits"]),
+        pytest.param([abs_path+"sample_flc_header.fits", abs_path+"sample_drc_header.fits"], marks=pytest.mark.xfail),
+        pytest.param(["bad_filename.fits"], marks=pytest.mark.xfail),
+    ],
+)
+def test_make_mosaic_wcs(filename_array):
+    new_wcs = make_mosaic_wcs(filename_array)

drizzlepac/wcs_functions.py

@@ -6,6 +6,7 @@
 
 """
 from astropy.io import fits as pyfits
+from astropy.wcs.utils import is_proj_plane_distorted
 import copy
 import numpy as np
 from numpy import linalg
@@ -129,8 +130,8 @@ def get_hstwcs(filename, extnum):
     """ Return the HSTWCS object for a given chip. """
     hdrwcs = wcsutil.HSTWCS(filename, ext=extnum)
     hdrwcs.filename = filename
-    hdrwcs.expname = pyfits.getval(filename, 'expname', ext=extnum)
-    hdrwcs.extver = pyfits.getval(filename, 'extver', ext=extnum)
+    # hdrwcs.expname = pyfits.getval(filename, 'expname', ext=extnum)
+    # hdrwcs.extver = pyfits.getval(filename, 'extver', ext=extnum)
 
     return hdrwcs
 
@@ -1069,15 +1070,17 @@ def make_mosaic_wcs(filenames, rot=None, scale=None):
     mosaic_wcs : HSTWCS object
         the merged composite WCS
     """
+
+    # makes str into list of str
     if not isinstance(filenames, list):
         filenames = [filenames]
 
-    # Compile list of WCSs for all chips from all input filenames
-    hstwcs_list = []
-    for f in filenames:
-        hstwcs_list.extend([get_hstwcs(f, extnum) for extnum in get_extns(f)])
-    # Combine them into a single mosaic WCS
-    output_wcs = utils.output_wcs(hstwcs_list, undistort=True)
+    hstwcs_list = [get_hstwcs(f, extnum) for f in filenames for extnum in get_extns(f)]
+
+    # Generate output WCS based on input WCSs, and undistorts depending on extension.
+    output_wcs = utils.output_wcs(hstwcs_list, undistort=is_proj_plane_distorted(hstwcs_list[0]))
+
+    # Combine them into a single mosaic WCS    
     output_wcs.wcs.cd = make_perfect_cd(output_wcs)
 
     # Apply each parameter in order, as this is effectively what is