dkist-processing-visp 3.3.0__py3-none-any.whl → 5.1.1__py3-none-any.whl

dkist_processing_visp/tests/test_parse.py

@@ -2,7 +2,6 @@ import pytest
 from dkist_data_simulator.dataset import key_function
 from dkist_processing_common._util.scratch import WorkflowFileSystem
 from dkist_processing_common.models.tags import Tag
-from dkist_processing_common.tests.conftest import FakeGQLClient
 
 from dkist_processing_visp.models.parameters import VispParsingParameters
 from dkist_processing_visp.models.tags import VispTag
@@ -87,6 +86,7 @@ def write_input_dark_frames_to_task(
     time_delta: float = 10.0,
     num_modstates: int = 2,
     data_shape: tuple[int, int] = (2, 2),
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = VispHeadersInputDarkFrames(
@@ -95,6 +95,7 @@ def write_input_dark_frames_to_task(
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
         num_modstates=num_modstates,
+        **kwargs,
     )
 
     num_written_frames = write_frames_to_task(
@@ -110,6 +111,7 @@ def write_input_lamp_frames_to_task(
     time_delta: float = 10.0,
     num_modstates: int = 2,
     data_shape: tuple[int, int] = (2, 2),
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = VispHeadersInputLampGainFrames(
@@ -118,6 +120,7 @@ def write_input_lamp_frames_to_task(
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
         num_modstates=num_modstates,
+        **kwargs,
     )
 
     num_written_frames = write_frames_to_task(
@@ -133,6 +136,7 @@ def write_input_solar_frames_to_task(
     time_delta: float = 10.0,
     num_modstates: int = 2,
     data_shape: tuple[int, int] = (2, 2),
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = VispHeadersInputSolarGainFrames(
@@ -141,6 +145,7 @@ def write_input_solar_frames_to_task(
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
         num_modstates=num_modstates,
+        **kwargs,
     )
 
     num_written_frames = write_frames_to_task(
@@ -156,6 +161,7 @@ def write_input_polcal_frames_to_task(
     time_delta: float = 30.0,
     num_modstates: int = 2,
     data_shape: tuple[int, int] = (2, 2),
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = VispHeadersInputPolcalFrames(
@@ -164,6 +170,7 @@ def write_input_polcal_frames_to_task(
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
         num_modstates=num_modstates,
+        **kwargs,
     )
 
     num_written_frames = write_frames_to_task(
@@ -179,6 +186,7 @@ def write_input_polcal_dark_frames_to_task(
     time_delta: float = 30.0,
     num_modstates: int = 2,
     data_shape: tuple[int, int] = (2, 2),
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = VispHeadersInputPolcalDarkFrames(
@@ -187,6 +195,7 @@ def write_input_polcal_dark_frames_to_task(
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
         num_modstates=num_modstates,
+        **kwargs,
     )
 
     num_written_frames = write_frames_to_task(
@@ -202,6 +211,7 @@ def write_input_polcal_gain_frames_to_task(
     time_delta: float = 30.0,
     num_modstates: int = 2,
     data_shape: tuple[int, int] = (2, 2),
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = VispHeadersInputPolcalGainFrames(
@@ -210,6 +220,7 @@ def write_input_polcal_gain_frames_to_task(
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
         num_modstates=num_modstates,
+        **kwargs,
     )
 
     num_written_frames = write_frames_to_task(
@@ -228,6 +239,7 @@ def write_input_observe_frames_to_task(
     time_delta: float = 10.0,
     data_shape: tuple[int, int] = (2, 2),
     obs_dataset_class=VispHeadersValidObserveFrames,
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = obs_dataset_class(
@@ -238,6 +250,7 @@ def write_input_observe_frames_to_task(
         num_modstates=num_modstates,
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
+        **kwargs,
     )
     num_written_frames = write_frames_to_task(
         task=task, frame_generator=dataset, extra_tags=[VispTag.input()]
@@ -273,6 +286,11 @@ def write_input_cal_frames_to_task(
     solar_exp_time,
     polcal_exp_time,
     num_modstates,
+    testing_arm_id,
+    testing_solar_ip_start_time,
+    testing_grating_constant,
+    testing_grating_angle,
+    testing_arm_position,
 ):
     def write_frames_to_task(task):
         for readout_exp_time in required_dark_readout_exp_times:
@@ -281,6 +299,7 @@ def write_input_cal_frames_to_task(
                 readout_exp_time=readout_exp_time,
                 exp_time=dark_exp_time,
                 num_modstates=num_modstates,
+                arm_id=testing_arm_id,
             )
 
         write_input_lamp_frames_to_task(
@@ -288,30 +307,39 @@ def write_input_cal_frames_to_task(
             readout_exp_time=lamp_readout_exp_time,
             exp_time=lamp_exp_time,
             num_modstates=num_modstates,
+            arm_id=testing_arm_id,
         )
         write_input_solar_frames_to_task(
             task=task,
             readout_exp_time=solar_readout_exp_time,
             exp_time=solar_exp_time,
             num_modstates=num_modstates,
+            arm_id=testing_arm_id,
+            ip_start_time=testing_solar_ip_start_time,
+            grating_constant=testing_grating_constant,
+            grating_angle=testing_grating_angle,
+            arm_position=testing_arm_position,
         )
         write_input_polcal_frames_to_task(
             task=task,
             readout_exp_time=polcal_readout_exp_time,
             exp_time=polcal_exp_time,
             num_modstates=num_modstates,
+            arm_id=testing_arm_id,
         )
         write_input_polcal_dark_frames_to_task(
             task=task,
             readout_exp_time=polcal_readout_exp_time,
             exp_time=polcal_exp_time,
             num_modstates=num_modstates,
+            arm_id=testing_arm_id,
         )
         write_input_polcal_gain_frames_to_task(
             task,
             readout_exp_time=polcal_readout_exp_time,
             exp_time=polcal_exp_time,
             num_modstates=num_modstates,
+            arm_id=testing_arm_id,
         )
 
     return write_frames_to_task
@@ -350,6 +378,11 @@ def test_parse_visp_input_data(
     observe_exp_times,
     num_modstates,
     mocker,
+    fake_gql_client,
+    testing_arm_id,
+    testing_grating_constant,
+    testing_grating_angle,
+    testing_arm_position,
 ):
     """
     Given: A ParseVispInputData task
@@ -357,7 +390,7 @@ def test_parse_visp_input_data(
     Then: All tagged files exist and individual task tags are applied
     """
     mocker.patch(
-        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
     )
     task = parse_task_with_no_data
     write_input_cal_frames_to_task(task)
@@ -369,6 +402,10 @@ def test_parse_visp_input_data(
             num_steps=3,
             readout_exp_time=obs_readout_exp_time,
             exp_time=obs_exp_time,
+            arm_id=testing_arm_id,
+            grating_constant=testing_grating_constant,
+            grating_angle=testing_grating_angle,
+            arm_position=testing_arm_position,
         )
 
     # When
@@ -393,6 +430,7 @@ def test_parse_visp_input_data_constants(
     parse_task_with_no_data,
     write_input_cal_frames_to_task,
     mocker,
+    fake_gql_client,
     lamp_readout_exp_time,
     solar_readout_exp_time,
     polcal_readout_exp_time,
@@ -403,6 +441,12 @@ def test_parse_visp_input_data_constants(
     polcal_exp_time,
     observe_exp_times,
     num_modstates,
+    testing_arm_id,
+    testing_obs_ip_start_time,
+    testing_solar_ip_start_time,
+    testing_grating_constant,
+    testing_grating_angle,
+    testing_arm_position,
 ):
     """
     Given: A ParseVispInputData task
@@ -410,7 +454,7 @@ def test_parse_visp_input_data_constants(
     Then: Constants are in the constants object as expected
     """
     mocker.patch(
-        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
     )
     task = parse_task_with_no_data
     write_input_cal_frames_to_task(task)
@@ -425,16 +469,29 @@ def test_parse_visp_input_data_constants(
             num_steps=num_steps,
             readout_exp_time=obs_readout_exp_time,
             exp_time=obs_exp_time,
+            arm_id=testing_arm_id,
+            ip_start_time=testing_obs_ip_start_time,
+            grating_constant=testing_grating_constant,
+            grating_angle=testing_grating_angle,
+            arm_position=testing_arm_position,
         )
 
     # When
     task()
     # Then
+    assert task.constants._db_dict["ARM_ID"] == testing_arm_id
     expected_dark_readout_exp_times = [
         lamp_readout_exp_time,
         solar_readout_exp_time,
     ] + observe_readout_exp_times
-    assert task.constants._db_dict["OBS_IP_START_TIME"] == "2022-11-28T13:55:00"
+    assert task.constants._db_dict["OBS_IP_START_TIME"] == testing_obs_ip_start_time
+    assert task.constants._db_dict["INCIDENT_LIGHT_ANGLE_DEG"] == -1 * testing_grating_angle
+    assert (
+        task.constants._db_dict["REFLECTED_LIGHT_ANGLE_DEG"]
+        == -1 * testing_grating_angle + testing_arm_position
+    )
+    assert task.constants._db_dict["GRATING_CONSTANT_INVERSE_MM"] == testing_grating_constant
+    assert task.constants._db_dict["SOLAR_GAIN_IP_START_TIME"] == testing_solar_ip_start_time
     assert task.constants._db_dict["NUM_MODSTATES"] == num_modstates
     assert task.constants._db_dict["NUM_MAP_SCANS"] == num_maps_per_readout_exp_time * len(
         observe_readout_exp_times
@@ -454,6 +511,10 @@ def test_parse_visp_input_data_constants(
         observe_readout_exp_times
     )
     assert task.constants._db_dict["RETARDER_NAME"] == "SiO2 OC"
+    assert task.constants._db_dict["DARK_GOS_LEVEL3_STATUS"] == "lamp"
+    assert task.constants._db_dict["SOLAR_GAIN_GOS_LEVEL3_STATUS"] == "clear"
+    assert task.constants._db_dict["SOLAR_GAIN_NUM_RAW_FRAMES_PER_FPA"] == 10
+    assert task.constants._db_dict["POLCAL_NUM_RAW_FRAMES_PER_FPA"] == 10
 
 
 def test_parse_visp_values(
@@ -462,6 +523,11 @@ def test_parse_visp_values(
     observe_readout_exp_times,
     num_modstates,
     mocker,
+    fake_gql_client,
+    testing_arm_id,
+    testing_grating_constant,
+    testing_grating_angle,
+    testing_arm_position,
 ):
     """
     :Given: A valid parse input task
@@ -469,7 +535,7 @@ def test_parse_visp_values(
     :Then: Values are correctly loaded into the constants mutable mapping
     """
     mocker.patch(
-        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
     )
     task = parse_task_with_no_data
     write_input_cal_frames_to_task(task)
@@ -480,6 +546,10 @@ def test_parse_visp_values(
         num_maps=1,
         num_steps=1,
         num_modstates=num_modstates,
+        arm_id=testing_arm_id,
+        grating_constant=testing_grating_constant,
+        grating_angle=testing_grating_angle,
+        arm_position=testing_arm_position,
     )
 
     task()
@@ -488,16 +558,19 @@ def test_parse_visp_values(
     assert task.constants.maximum_cadence == 10
     assert task.constants.minimum_cadence == 10
     assert task.constants.variance_cadence == 0
+    assert task.constants.camera_name == "camera_name"
 
 
-def test_multiple_num_raster_steps_raises_error(parse_task_with_no_data, num_modstates, mocker):
+def test_multiple_num_raster_steps_raises_error(
+    parse_task_with_no_data, num_modstates, mocker, fake_gql_client
+):
     """
     :Given: A prase task with data that have inconsistent VSPNSTP values
     :When: Calling the parse task
     :Then: The correct error is raised
     """
     mocker.patch(
-        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
     )
     task = parse_task_with_no_data
     write_input_dark_frames_to_task(task, readout_exp_time=0.1, exp_time=0.2)
@@ -515,14 +588,14 @@ def test_multiple_num_raster_steps_raises_error(parse_task_with_no_data, num_mod
         task()
 
 
-def test_incomplete_single_map(parse_task_with_no_data, num_modstates, mocker):
+def test_incomplete_single_map(parse_task_with_no_data, num_modstates, mocker, fake_gql_client):
     """
     :Given: A parse task with data that has an incomplete raster scan
     :When: Calling the parse task
     :Then: The correct number of raster steps are found
     """
     mocker.patch(
-        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
     )
     task = parse_task_with_no_data
     num_steps = 4
@@ -542,14 +615,14 @@ def test_incomplete_single_map(parse_task_with_no_data, num_modstates, mocker):
     assert task.constants._db_dict["NUM_MAP_SCANS"] == num_map_scans
 
 
-def test_incomplete_final_map(parse_task_with_no_data, num_modstates, mocker):
+def test_incomplete_final_map(parse_task_with_no_data, num_modstates, mocker, fake_gql_client):
     """
     :Given: A parse task with data that has complete raster scans along with an incomplete raster scan
     :When: Calling the parse task
     :Then: The correct number of raster steps and maps are found
     """
     mocker.patch(
-        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
     )
     task = parse_task_with_no_data
     num_steps = 4
@@ -575,6 +648,11 @@ def test_intensity_observes_and_polarimetric_cals(
     observe_readout_exp_times,
     observe_exp_times,
     mocker,
+    fake_gql_client,
+    testing_arm_id,
+    testing_grating_constant,
+    testing_grating_angle,
+    testing_arm_position,
 ):
     """
     :Given: Data where the observe frames are in intensity mode and the calibration frames are in polarimetric mode
@@ -582,7 +660,7 @@ def test_intensity_observes_and_polarimetric_cals(
     :Then: All modulator state keys generated for all frames are in the first modulator state
     """
     mocker.patch(
-        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
     )
     task = parse_task_with_no_data
     write_input_cal_frames_to_task(task)
@@ -594,6 +672,10 @@ def test_intensity_observes_and_polarimetric_cals(
         readout_exp_time=observe_readout_exp_times[0],
         exp_time=observe_exp_times[0],
         obs_dataset_class=VispHeadersIntensityObserveFrames,
+        arm_id=testing_arm_id,
+        grating_constant=testing_grating_constant,
+        grating_angle=testing_grating_angle,
+        arm_position=testing_arm_position,
     )
     task()
     assert task.constants._db_dict["NUM_MODSTATES"] == 1
@@ -603,14 +685,16 @@ def test_intensity_observes_and_polarimetric_cals(
         assert "MODSTATE_1" in task.scratch.tags(file)
 
 
-def test_dark_readout_exp_time_picky_bud(parse_task_with_no_data, mocker, lamp_readout_exp_time):
+def test_dark_readout_exp_time_picky_bud(
+    parse_task_with_no_data, mocker, fake_gql_client, lamp_readout_exp_time
+):
     """
     :Given: Dataset where non-dark readout exp time values are missing from the set of dark IP frames.
     :When: Parsing
     :Then: The `DarkReadoutExpTimePickyBud` raises an error
     """
     mocker.patch(
-        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
     )
     bad_readout_exp_time = lamp_readout_exp_time + 0.02
     dummy_exp_time = 99.0

dkist_processing_visp/tests/test_quality.py

@@ -6,7 +6,6 @@ from dkist_processing_common._util.scratch import WorkflowFileSystem
 from dkist_processing_common.codecs.fits import fits_array_encoder
 from dkist_processing_common.models.tags import Tag
 from dkist_processing_common.models.task_name import TaskName
-from dkist_processing_common.tests.conftest import FakeGQLClient
 
 from dkist_processing_visp.models.tags import VispTag
 from dkist_processing_visp.tasks.quality_metrics import VispL0QualityMetrics
@@ -132,7 +131,7 @@ def test_l0_quality_task(
 
 
 @pytest.mark.parametrize("pol_mode", ["observe_polarimetric", "observe_intensity"])
-def test_l1_quality_task(visp_l1_quality_task, pol_mode, mocker):
+def test_l1_quality_task(visp_l1_quality_task, pol_mode, mocker, fake_gql_client):
     """
     Given: A VispL1QualityMetrics task
     When: Calling the task instance
@@ -140,7 +139,7 @@ def test_l1_quality_task(visp_l1_quality_task, pol_mode, mocker):
             and a single noise measurement and datetime is recorded for L1 file for each Stokes Q, U, and V
     """
     mocker.patch(
-        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
     )
     # When
     task, num_maps, num_steps, num_stokes = visp_l1_quality_task

dkist_processing_visp/tests/test_science.py

@@ -12,8 +12,8 @@ from dkist_header_validator import spec122_validator
 from dkist_processing_common._util.scratch import WorkflowFileSystem
 from dkist_processing_common.codecs.fits import fits_array_encoder
 from dkist_processing_common.codecs.fits import fits_hdu_decoder
+from dkist_processing_common.models.fits_access import MetadataKey
 from dkist_processing_common.models.tags import Tag
-from dkist_processing_common.tests.conftest import FakeGQLClient
 
 from dkist_processing_visp.models.tags import VispStemName
 from dkist_processing_visp.models.tags import VispTag
@@ -144,7 +144,7 @@ def dummy_calibration_collection():
 
     dark_dict = {VispTag.beam(beam): {VispTag.readout_exp_time(0.04): np.zeros(intermediate_shape)}}
     background_dict = {VispTag.beam(beam): np.zeros(intermediate_shape)}
-    solar_dict = {VispTag.beam(beam): {VispTag.modstate(modstate): np.ones(intermediate_shape)}}
+    solar_dict = {VispTag.beam(beam): np.ones(intermediate_shape)}
     angle_dict = {VispTag.beam(beam): 0.0}
     spec_dict = {VispTag.beam(beam): np.zeros(intermediate_shape[1])}
     offset_dict = {VispTag.beam(beam): {VispTag.modstate(modstate): np.zeros(2)}}
@@ -184,7 +184,7 @@ def headers_with_dates() -> tuple[list[fits.Header], str, int, int]:
     ]
     random.shuffle(headers)  # Shuffle to make sure they're not already in time order
     for h in headers:
-        h["XPOSURE"] = exp_time  # Exposure time, in ms
+        h[MetadataKey.fpa_exposure_time_ms] = exp_time  # Exposure time, in ms
 
     return headers, start_time, exp_time, time_delta
 
@@ -252,10 +252,14 @@ def calibration_collection_with_full_overlap_slice() -> CalibrationCollection:
 
 @pytest.mark.parametrize(
     "background_on",
-    [pytest.param(True, id="Background on"), pytest.param(False, id="Background off")],
+    [pytest.param(True, id="background_on"), pytest.param(False, id="background_off")],
 )
 def test_science_calibration_task(
-    science_calibration_task, background_on, assign_input_dataset_doc_to_task, mocker
+    science_calibration_task,
+    background_on,
+    assign_input_dataset_doc_to_task,
+    mocker,
+    fake_gql_client,
 ):
     """
     Given: A ScienceCalibration task
@@ -264,7 +268,7 @@ def test_science_calibration_task(
     """
 
     mocker.patch(
-        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
     )
 
     # When
@@ -304,7 +308,8 @@ def test_science_calibration_task(
         task=task, num_modstates=num_modstates, data_shape=intermediate_shape, offsets=offsets
     )
     write_dummy_intermediate_solar_cals_to_task(
-        task=task, data_shape=intermediate_shape, num_modstates=num_modstates
+        task=task,
+        data_shape=intermediate_shape,
     )
     write_demod_matrices_to_task(task=task, num_modstates=num_modstates)
     write_input_observe_frames_to_task(
@@ -360,10 +365,10 @@ def test_science_calibration_task(
         assert header["VSPMAP"] == map_scan
 
         # Check that WCS keys were updated
-        if offsets[1, 0, 0] > 0:
-            assert header["CRPIX2"] == input_header["CRPIX2"] - np.ceil(offsets[1, 0, 0])
-        if offsets[1, 0, 1] > 0:
-            assert header["CRPIX1"] == input_header["CRPIX1"] - np.ceil(offsets[1, 0, 1])
+        if offsets[1, 0, 0] < 0:
+            assert header["CRPIX2"] == input_header["CRPIX2"] - np.ceil(-offsets[1, 0, 0])
+        if offsets[1, 0, 1] < 0:
+            assert header["CRPIX1"] == input_header["CRPIX1"] - np.ceil(-offsets[1, 0, 1])
 
     quality_files = task.read(tags=[Tag.quality("TASK_TYPES")])
     for file in quality_files:
@@ -416,7 +421,7 @@ def test_readout_normalization_correct(
     )
 
     # When:
-    corrected_array, _ = task.correct_single_frame(
+    corrected_array, _, _ = task.correct_single_frame(
         beam=1,
         modstate=1,
         raster_step=1,
@@ -509,7 +514,7 @@ def test_compute_date_keys_compressed_headers(
                     [[1.0, 2.0], [11.0, 10.0], [3.0, 2.0]],  # Beam 2
                 ]
             ),
-            [slice(11, None, None), slice(10, None, None)],
+            [slice(0, -11, None), slice(0, -10, None)],
         ),
         (
             np.array(
@@ -518,7 +523,7 @@ def test_compute_date_keys_compressed_headers(
                     [[-1.0, -2.0], [-11.0, -10.0], [-3.0, -2.0]],  # Beam 2
                 ]
             ),
-            [slice(0, -11, None), slice(0, -10, None)],
+            [slice(11, None, None), slice(10, None, None)],
         ),
         (
             np.array(
@@ -527,7 +532,7 @@ def test_compute_date_keys_compressed_headers(
                     [[1.0, 2.0], [-11.0, 10.0], [-3.0, -2.0]],  # Beam 2
                 ]
             ),
-            [slice(10, -11, None), slice(10, -2, None)],
+            [slice(11, -10, None), slice(2, -10, None)],
         ),
     ],
     ids=["All positive", "All negative", "Positive and negative"],
@@ -577,3 +582,96 @@ def test_combine_beams(
         expected = np.ones((10, 10, 4)) * 2.5
 
     np.testing.assert_array_equal(data, expected)
+
+
+@pytest.mark.parametrize(
+    "shifts",
+    # Shifts have shape (num_beams, num_modstates, 2)
+    # So the inner-most lists below (e.g., [5.0, 6.0]) correspond to [x_shift, y_shit]
+    [
+        np.array(
+            [
+                [[0.0, 0.0], [10.0, 2.0], [5.0, 6.0]],  # Beam 1
+                [[1.0, 2.0], [-11.0, 10.0], [-3.0, -2.0]],  # Beam 2
+            ]
+        ),
+    ],
+    ids=["Positive and negative"],
+)
+def test_combine_and_cut_nan_masks(
+    science_calibration_task, calibration_collection_with_geo_shifts, shifts
+):
+    """
+    Given: A ScienceCalibration task and NaN masks, along with geometric shifts
+    When: Combining the two NaN masks
+    Then: The final mask has NaN values in the correct place and is correctly cropped
+    """
+    nan_1_location = [0, 1]
+    nan_2_location = [50, 50]
+    nan_3_location = [4, 1]
+    nan_4_location = [55, 63]
+    nan_mask_shape = (100, 100)
+    nan_mask_1 = np.zeros(shape=nan_mask_shape)
+    nan_mask_1[nan_1_location[0], nan_1_location[1]] = np.nan
+    nan_mask_1[nan_2_location[0], nan_2_location[1]] = np.nan
+    nan_mask_2 = np.zeros(shape=nan_mask_shape)
+    nan_mask_2[nan_3_location[0], nan_3_location[1]] = np.nan
+    nan_mask_2[nan_4_location[0], nan_4_location[1]] = np.nan
+    task, _, _, _, _, _ = science_calibration_task
+    combined_nan_mask = task.combine_and_cut_nan_masks(
+        nan_masks=[nan_mask_1, nan_mask_2], calibrations=calibration_collection_with_geo_shifts
+    )
+    beam_1_shifts = shifts[0]
+    beam_2_shifts = shifts[1]
+    beam_1_x_shifts = [i[0] for i in beam_1_shifts]
+    beam_2_x_shifts = [i[0] for i in beam_2_shifts]
+    beam_1_y_shifts = [i[1] for i in beam_1_shifts]
+    beam_2_y_shifts = [i[1] for i in beam_2_shifts]
+    x_shifts = beam_1_x_shifts + beam_2_x_shifts
+    y_shifts = beam_1_y_shifts + beam_2_y_shifts
+    assert combined_nan_mask.shape == (
+        nan_mask_shape[0] - (max(x_shifts) - min(x_shifts)),
+        nan_mask_shape[1] - (max(y_shifts) - min(y_shifts)),
+    )
+    # Check that one NaN value from each original mask is present in the combined mask and in the correct place
+    assert (
+        combined_nan_mask[
+            nan_2_location[0] - int(abs(min(x_shifts))), nan_2_location[1] - int(abs(min(y_shifts)))
+        ]
+        == True
+    )
+    assert (
+        combined_nan_mask[
+            nan_4_location[0] - int(abs(min(x_shifts))), nan_4_location[1] - int(abs(min(y_shifts)))
+        ]
+        == True
+    )
+    assert np.sum(combined_nan_mask) == 2  # only two NaN values are in the final mask
+
+
+def test_generate_nan_mask(science_calibration_task, dummy_calibration_collection):
+    """
+    Given: a calibration collection
+    When: calculating the NaN mask to use
+    Then: the mask takes up some, but not all, of the frame size
+    """
+    task, _, _, _, _, _ = science_calibration_task
+    calibration_collection, _, _ = dummy_calibration_collection
+    beam = 1
+    modstate = 1
+    solar_gain_array = calibration_collection.solar_gain[VispTag.beam(beam)]
+    angle = calibration_collection.angle[VispTag.beam(beam)]
+    spec_shift = calibration_collection.spec_shift[VispTag.beam(beam)]
+    state_offset = calibration_collection.state_offset[VispTag.beam(beam)][
+        VispTag.modstate(modstate)
+    ]
+    nan_mask = task.generate_nan_mask(
+        solar_corrected_array=np.random.random(size=solar_gain_array.shape),
+        state_offset=state_offset,
+        angle=angle,
+        spec_shift=spec_shift,
+    )
+    # Some of the mask is marked as NaN but not all
+    assert np.sum(nan_mask) < np.size(nan_mask)
+    # Ensure that only zeroes and ones are in the mask
+    assert set(np.unique(nan_mask)) == {0, 1}