dkist-processing-common 10.5.4__py3-none-any.whl → 12.1.0rc1__py3-none-any.whl

dkist_processing_common/tests/test_write_l1.py

@@ -18,13 +18,70 @@ from dkist_processing_common._util.scratch import WorkflowFileSystem
 from dkist_processing_common.codecs.fits import fits_hdu_decoder
 from dkist_processing_common.codecs.fits import fits_hdulist_encoder
 from dkist_processing_common.models.graphql import RecipeRunProvenanceResponse
-from dkist_processing_common.models.graphql import RecipeRunResponse
 from dkist_processing_common.models.tags import Tag
 from dkist_processing_common.models.wavelength import WavelengthRange
 from dkist_processing_common.tasks.write_l1 import WriteL1Frame
-from dkist_processing_common.tests.conftest import FakeGQLClient
-from dkist_processing_common.tests.conftest import FakeGQLClientNoRecipeConfiguration
-from dkist_processing_common.tests.conftest import TILE_SIZE
+from dkist_processing_common.tests.mock_metadata_store import TILE_SIZE
+from dkist_processing_common.tests.mock_metadata_store import InputDatasetRecipeRunResponseMapping
+from dkist_processing_common.tests.mock_metadata_store import RecipeRunResponseMapping
+from dkist_processing_common.tests.mock_metadata_store import fake_gql_client_factory
+from dkist_processing_common.tests.mock_metadata_store import (
+    make_default_input_dataset_recipe_run_response,
+)
+from dkist_processing_common.tests.mock_metadata_store import make_default_recipe_run_response
+
+
+@pytest.fixture
+def fake_gql_client_default_configuration():
+    """Create GraphQL client Mock that returns result without recipe run configuration."""
+    recipe_run_response = make_default_recipe_run_response()
+    recipe_run_response.configuration = None
+    new_response_mapping = RecipeRunResponseMapping(response=recipe_run_response)
+    FakeGQLClientDefaultConfiguration = fake_gql_client_factory(
+        response_mapping_override=new_response_mapping
+    )
+
+    return FakeGQLClientDefaultConfiguration
+
+
+@pytest.fixture
+def fake_gql_client_missing_calibration_part():
+    """Create GraphQL client Mock that returns result without calibration part."""
+    input_dataset_recipe_run_response = make_default_input_dataset_recipe_run_response()
+    dataset_parts = (
+        input_dataset_recipe_run_response.recipeInstance.inputDataset.inputDatasetInputDatasetParts
+    )
+    for index, part in enumerate(dataset_parts):
+        if (
+            part.inputDatasetPart.inputDatasetPartType.inputDatasetPartTypeName
+            == "calibration_frames"
+        ):
+            del dataset_parts[index]
+    new_response_mapping = InputDatasetRecipeRunResponseMapping(
+        response=input_dataset_recipe_run_response
+    )
+    FakeGQLClientMissingInputDatasetCalibrationPart = fake_gql_client_factory(
+        response_mapping_override=new_response_mapping
+    )
+
+    return FakeGQLClientMissingInputDatasetCalibrationPart
+
+
+@pytest.fixture()
+def make_fake_gql_client_with_provenance():
+    """Create GraphQL client Mocks that will return customizable provenance records."""
+
+    def class_generator(provenances: list[RecipeRunProvenanceResponse]):
+        recipe_run_response = make_default_recipe_run_response()
+        recipe_run_response.recipeRunProvenances = provenances
+        new_response_mapping = RecipeRunResponseMapping(response=recipe_run_response)
+        FakeGQLClientProvenances = fake_gql_client_factory(
+            response_mapping_override=new_response_mapping
+        )
+
+        return FakeGQLClientProvenances
+
+    return class_generator
 
 
 class CompleteWriteL1Frame(WriteL1Frame):
@@ -87,6 +144,12 @@ class CompleteWriteL1Frame(WriteL1Frame):
         return WavelengthRange(min=1075.0 * u.nm, max=1085.0 * u.nm)
 
 
+class CompleteWriteL1FrameWithEmptyWaveband(CompleteWriteL1Frame):
+    def get_wavelength_range(self, header: fits.Header) -> WavelengthRange:
+        # Return an empty range to test the empty waveband case
+        return WavelengthRange(min=10000.0 * u.nm, max=10050.0 * u.nm)
+
+
 @dataclass
 class FakeConstantDb:
     INSTRUMENT: str = "TEST"
@@ -121,6 +184,45 @@ def write_l1_task(request, recipe_run_id, tmp_path):
         stokes_params = ["I", "Q", "U", "V"]
         used_stokes_params = []
         hdu = fits.PrimaryHDU(data=np.random.random(size=(1, 128, 128)) * 10, header=header)
+        hdu.header["IPTASK"] = "level0_only key to be removed"
+        hdul = fits.HDUList([hdu])
+        for i in range(num_of_stokes_params):
+            task.write(
+                data=hdul,
+                tags=[
+                    Tag.calibrated(),
+                    Tag.frame(),
+                    Tag.stokes(stokes_params[i]),
+                    Tag.dsps_repeat(i),
+                ],
+                encoder=fits_hdulist_encoder,
+            )
+            used_stokes_params.append(stokes_params[i])
+        task.constants._update(asdict(FakeConstantDb()))
+        yield task, used_stokes_params, header
+        task._purge()
+
+
+@pytest.fixture(
+    scope="function",
+    params=[
+        pytest.param((1, "complete_common_header"), id="Intensity"),
+        pytest.param((4, "complete_polarimetric_header"), id="Polarimetric"),
+    ],
+)
+def write_l1_task_with_empty_waveband(recipe_run_id, tmp_path, request):
+    with CompleteWriteL1FrameWithEmptyWaveband(
+        recipe_run_id=recipe_run_id,
+        workflow_name="workflow_name",
+        workflow_version="workflow_version",
+    ) as task:
+        task.scratch = WorkflowFileSystem(recipe_run_id=recipe_run_id, scratch_base_path=tmp_path)
+        num_of_stokes_params, header_fixture_name = request.param
+        header = request.getfixturevalue(header_fixture_name)
+        stokes_params = ["I", "Q", "U", "V"]
+        used_stokes_params = []
+        hdu = fits.PrimaryHDU(data=np.random.random(size=(1, 128, 128)) * 10, header=header)
+        hdu.header["IPTASK"] = "level0_only key to be removed"
         hdul = fits.HDUList([hdu])
         for i in range(num_of_stokes_params):
             task.write(
@@ -142,23 +244,46 @@ def write_l1_task(request, recipe_run_id, tmp_path):
 @pytest.fixture(
     scope="function",
     params=[
-        pytest.param(True, id="AO_lock_on"),
-        pytest.param(False, id="AO_lock_off"),
+        pytest.param(
+            {"AO_LOCK": True, "ATMOS_R0": 0.2, "OOBSHIFT": 17}, id="AO_LOCK_True_good_R0_good_oob"
+        ),
+        pytest.param(
+            {"AO_LOCK": True, "ATMOS_R0": 1, "OOBSHIFT": 17}, id="AO_LOCK_True_bad_R0_good_oob"
+        ),
+        pytest.param(
+            {"AO_LOCK": False, "ATMOS_R0": 0.2, "OOBSHIFT": 17}, id="AO_LOCK_False_good_R0_good_oob"
+        ),
+        pytest.param(
+            {"AO_LOCK": False, "ATMOS_R0": 1, "OOBSHIFT": 17}, id="AO_LOCK_False_bad_R0_good_oob"
+        ),
+        pytest.param(
+            {"AO_LOCK": True, "ATMOS_R0": 0.2, "OOBSHIFT": 150}, id="AO_LOCK_True_good_R0_bad_oob"
+        ),
+        pytest.param(
+            {"AO_LOCK": True, "ATMOS_R0": 1, "OOBSHIFT": 150}, id="AO_LOCK_True_bad_R0_bad_oob"
+        ),
+        pytest.param(
+            {"AO_LOCK": False, "ATMOS_R0": 0.2, "OOBSHIFT": 150}, id="AO_LOCK_False_good_R0_bad_oob"
+        ),
+        pytest.param(
+            {"AO_LOCK": False, "ATMOS_R0": 1, "OOBSHIFT": 150}, id="AO_LOCK_False_bad_R0_bad_oob"
+        ),
+        pytest.param({"ATMOS_R0": 0.2, "OOBSHIFT": 17}, id="AO_LOCK_missing"),
+        pytest.param({"ATMOS_R0": 0.2, "AO_LOCK": True}, id="OOBSHIFT_missing"),
     ],
 )
 def write_l1_task_no_data(request, recipe_run_id, tmp_path, complete_common_header):
-    with (
-        CompleteWriteL1Frame(
-            recipe_run_id=recipe_run_id,
-            workflow_name="workflow_name",
-            workflow_version="workflow_version",
-        ) as task
-    ):
+    with CompleteWriteL1Frame(
+        recipe_run_id=recipe_run_id,
+        workflow_name="workflow_name",
+        workflow_version="workflow_version",
+    ) as task:
         task.scratch = WorkflowFileSystem(recipe_run_id=recipe_run_id, scratch_base_path=tmp_path)
         header = complete_common_header
-        header["AO_LOCK"] = request.param
-        fried_parameter = 0.2
-        header["ATMOS_R0"] = fried_parameter
+        header.pop("AO_LOCK", None)
+        header.pop("ATMOS_R0", None)
+        header.pop("OOBSHIFT", None)
+        header.update(request.param)
         hdu = fits.PrimaryHDU(data=np.random.random(size=(1, 1, 1)) * 1, header=header)
         hdul = fits.HDUList([hdu])
         task.write(
@@ -170,29 +295,12 @@ def write_l1_task_no_data(request, recipe_run_id, tmp_path, complete_common_head
             encoder=fits_hdulist_encoder,
         )
         task.constants._update(asdict(FakeConstantDb()))
-        yield task, header, fried_parameter
+        fried_parameter = request.param["ATMOS_R0"]
+        oob_shift = request.param.get("OOBSHIFT")
+        yield task, header, fried_parameter, oob_shift
         task._purge()
 
 
-@pytest.fixture()
-def make_mock_gql_client_with_provenance():
-    """Factory to create GraphQL client Mocks that will return customizable provenance records."""
-
-    def factory(provenances: list[RecipeRunProvenanceResponse]):
-        class WriteL1FakeGQLClient(FakeGQLClient):
-            def execute_gql_query(self, **kwargs):
-                response = super().execute_gql_query(**kwargs)
-                if isinstance(response, list):
-                    if isinstance(response[0], RecipeRunResponse):
-                        response: list[RecipeRunResponse]
-                        response[0].recipeRunProvenances = provenances
-                return response
-
-        return WriteL1FakeGQLClient
-
-    return factory
-
-
 @pytest.mark.parametrize(
     "provenances, is_manual",
     [
@@ -227,7 +335,7 @@ def make_mock_gql_client_with_provenance():
 def test_write_l1_frame(
     write_l1_task,
     mocker,
-    make_mock_gql_client_with_provenance,
+    make_fake_gql_client_with_provenance,
     provenances: list[RecipeRunProvenanceResponse],
     is_manual,
 ):
@@ -236,7 +344,7 @@ def test_write_l1_frame(
     :When: running the task
     :Then: no errors are raised and the MANPROC and FRAMEVOL headers are correct
     """
-    WriteL1GQLClient = make_mock_gql_client_with_provenance(provenances=provenances)
+    WriteL1GQLClient = make_fake_gql_client_with_provenance(provenances=provenances)
 
     mocker.patch(
         "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=WriteL1GQLClient
@@ -265,6 +373,9 @@ def test_write_l1_frame(
             # Test that FRAMEVOL still has its comment
             assert hdu.header.comments["FRAMEVOL"]
 
+            # Test that 'level0_only' keys are being removed
+            assert "IPTASK" not in hdu.header.keys()
+
 
 def test_replace_header_values(write_l1_task):
     """
@@ -316,14 +427,14 @@ def test_calculate_telapse(write_l1_task):
     assert task.calculate_telapse(header=header) == 86400
 
 
-def test_solarnet_keys(write_l1_task, mocker):
+def test_solarnet_keys(write_l1_task, mocker, fake_gql_client):
     """
     :Given: files with headers converted to SPEC 214 L1
     :When: checking the solarnet extra headers
     :Then: the correct values 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
     )
     mocker.patch(
         "dkist_processing_common.tasks.write_l1.WriteL1Frame.version_from_module_name",
@@ -348,14 +459,14 @@ def test_solarnet_keys(write_l1_task, mocker):
         assert header["WAVEMAX"] == 1085.0
 
 
-def test_documentation_keys(write_l1_task, mocker):
+def test_documentation_keys(write_l1_task, mocker, fake_gql_client):
     """
     :Given: files with headers converted to SPEC 214 L1
     :When: checking the documentation header URLs
     :Then: the correct values 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
     )
     mocker.patch(
         "dkist_processing_common.tasks.write_l1.WriteL1Frame.version_from_module_name",
@@ -386,9 +497,9 @@ def test_get_version_from_module(write_l1_task):
     assert task.version_from_module_name() == common_version
 
 
-def test_get_tile_size(write_l1_task, mocker):
+def test_get_tile_size(write_l1_task, mocker, fake_gql_client):
     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, _, _ = write_l1_task
     test_array = np.zeros((1, TILE_SIZE // 2, TILE_SIZE * 2))
@@ -396,28 +507,31 @@ def test_get_tile_size(write_l1_task, mocker):
     assert tile_size == [1, TILE_SIZE // 2, TILE_SIZE]
 
 
-def test_rice_compression_with_specified_tile_size(write_l1_task, mocker):
+def test_rice_compression_with_specified_tile_size(write_l1_task, mocker, fake_gql_client):
     """
     :Given: a write_L1 task with a specified tile size in the recipe configuration
     :When: running the task
     :Then: data is written with the compression tile size specified in the recipe configuration
     """
     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, _, _ = write_l1_task
     task()
     files = list(task.read(tags=[Tag.frame(), Tag.output()]))
     for file in files:
         hdul = fits.open(file)
-        comp_header = hdul[1]._header
+        bintable = hdul[1]._get_bintable_without_data()
+        comp_header = bintable.header
         data_shape = list(hdul[1].data.shape)
         data_shape.reverse()
         for i, dim in enumerate(data_shape):
             assert comp_header["ZTILE" + str(i + 1)] == min(dim, TILE_SIZE)
 
 
-def test_rice_compression_with_default_tile_size(write_l1_task, mocker):
+def test_rice_compression_with_default_tile_size(
+    write_l1_task, mocker, fake_gql_client_default_configuration
+):
     """
     :Given: a write_L1 task with no specified tile size in the recipe configuration
     :When: running the task
@@ -427,14 +541,16 @@ def test_rice_compression_with_default_tile_size(write_l1_task, mocker):
     """
     mocker.patch(
         "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient",
-        new=FakeGQLClientNoRecipeConfiguration,
+        new=fake_gql_client_default_configuration,
     )
     task, _, _ = write_l1_task
     task()
+    assert task.tile_size_param == None
     files = list(task.read(tags=[Tag.frame(), Tag.output()]))
     for file in files:
         hdul = fits.open(file)
-        comp_header = hdul[1]._header
+        bintable = hdul[1]._get_bintable_without_data()
+        comp_header = bintable.header
         data_shape = list(hdul[1].data.shape)
         data_shape.reverse()
         assert comp_header["ZTILE1"] == data_shape[0]
@@ -442,7 +558,7 @@ def test_rice_compression_with_default_tile_size(write_l1_task, mocker):
         assert comp_header["ZTILE3"] == 1
 
 
-def test_reprocessing_keys(write_l1_task, mocker):
+def test_reprocessing_keys(write_l1_task, mocker, fake_gql_client):
     """
     :Given: a write_L1 task with reprocessing keys present
     :When: running the task
@@ -450,18 +566,73 @@ def test_reprocessing_keys(write_l1_task, mocker):
     """
     mocker.patch(
         "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient",
-        new=FakeGQLClient,
+        new=fake_gql_client,
     )
     task, _, _ = write_l1_task
     task()
     files = list(task.read(tags=[Tag.frame(), Tag.output()]))
     for file in files:
         header = fits.open(file)[1].header
-        assert header["IDSPARID"] == task.metadata_store_input_dataset_parameters_part_id
-        assert header["IDSOBSID"] == task.metadata_store_input_dataset_observe_frames_part_id
-        assert header["IDSCALID"] == task.metadata_store_input_dataset_calibration_frames_part_id
+        assert header["IDSPARID"] == task.metadata_store_input_dataset_parameters.inputDatasetPartId
+        assert (
+            header["IDSOBSID"]
+            == task.metadata_store_input_dataset_observe_frames.inputDatasetPartId
+        )
+        assert (
+            header["IDSCALID"]
+            == task.metadata_store_input_dataset_calibration_frames.inputDatasetPartId
+        )
         assert header["WKFLNAME"] == task.workflow_name
         assert header["WKFLVERS"] == task.workflow_version
+        assert header["PROCTYPE"] == "L1"
+        assert header["PRODUCT"] == task.compute_product_id(header["IDSOBSID"], header["PROCTYPE"])
+
+
+def test_missing_input_dataset_part(
+    write_l1_task, mocker, fake_gql_client_missing_calibration_part
+):
+    """
+    :Given: a Write_L1 task with a missing calibration frames part
+    :When: running the task
+    :Then: the input dataset part keys are correctly written without throwing an exception
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient",
+        new=fake_gql_client_missing_calibration_part,
+    )
+    task, _, _ = write_l1_task
+    task()
+    files = list(task.read(tags=[Tag.frame(), Tag.output()]))
+    for file in files:
+        header = fits.open(file)[1].header
+        assert header["IDSPARID"] == task.metadata_store_input_dataset_parameters.inputDatasetPartId
+        assert (
+            header["IDSOBSID"]
+            == task.metadata_store_input_dataset_observe_frames.inputDatasetPartId
+        )
+        assert "IDSCALID" not in header
+
+
+@pytest.mark.parametrize(
+    "ids_obs_id, proc_type",
+    [
+        pytest.param(42, "alpha", id="42"),
+        pytest.param(1_000, "beta", id="thousand"),
+        pytest.param(1_000_000, "gamma", id="million"),
+    ],
+)
+def test_product_id_calculation(ids_obs_id: int, proc_type: str):
+    """
+    Given: integer IDSOBSID and string PROCTYPE
+    When: calculating the productId
+    Then: the productId is computed properly
+    """
+    product_id = WriteL1Frame.compute_product_id(ids_obs_id, proc_type)
+    assert isinstance(product_id, str)
+    assert product_id.startswith(f"{proc_type}-")
+    assert len(product_id) >= len(proc_type) + 6
+    # same result the second time around
+    assert product_id == WriteL1Frame.compute_product_id(ids_obs_id, proc_type)
 
 
 def test_calculate_date_end(write_l1_task):
@@ -493,7 +664,7 @@ def test_add_contributing_id_headers(write_l1_task):
     assert header["NEXPERS"] == 3
 
 
-def test_spectral_line_keys(write_l1_task, mocker):
+def test_spectral_line_keys(write_l1_task, mocker, fake_gql_client):
     """
     :Given: a header
     :When: adding spectral line information to the headers
@@ -501,10 +672,10 @@ def test_spectral_line_keys(write_l1_task, mocker):
     """
     mocker.patch(
         "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient",
-        new=FakeGQLClient,
+        new=fake_gql_client,
     )
     task, _, header = write_l1_task
-    header = task.add_datacenter_headers(header=header, hdu_size=1024, stokes="I")
+    header = task.add_spectral_line_headers(header=header)
     assert header["SPECLN01"] == "Fe XIII (1079.8 nm)"
     assert header["SPECLN02"] == "He I (1083.0 nm)"
     assert header["NSPECLNS"] == 2
@@ -512,32 +683,83 @@ def test_spectral_line_keys(write_l1_task, mocker):
         assert header["SPECLN03"]
 
 
-@pytest.mark.flaky(max_reruns=10)
-def test_location_of_dkist(write_l1_task):
-    """
-    Given: function for retrieving the dkist location on earth
-    When: Call function
-    Then: result is the same as what is in the astropy online database
-    """
-    task, _, _ = write_l1_task
-    itrs = task.location_of_dkist
-    assert itrs == EarthLocation.of_site("dkist")
-
-
 def test_check_r0_ao_lock(write_l1_task_no_data):
     """
     :Given: a header
     :When: writing, check if the AO lock is on
     :Then: write the r0 value if AO lock on, don't write if AO lock off
     """
-    task, header, r0 = write_l1_task_no_data
+    task, header, r0, _ = write_l1_task_no_data
     header_after_check = task.remove_invalid_r0_values(header=header)
-    if header["AO_LOCK"]:
+    if header.get("AO_LOCK"):
         assert header_after_check["ATMOS_R0"] == header["ATMOS_R0"]
         assert header["ATMOS_R0"] == r0
         assert header["AO_LOCK"]
-    if not header["AO_LOCK"]:
-        with pytest.raises(KeyError) as ke:
+    else:
+        with pytest.raises(KeyError, match="Keyword 'ATMOS_R0' not found"):
             invalid_r0 = header_after_check["ATMOS_R0"]
-        assert "Keyword 'ATMOS_R0' not found" in str(ke)
-        assert not header["AO_LOCK"]
+        assert header.get("AO_LOCK") != True
+
+
+@pytest.mark.parametrize(
+    "wavelength, wavemin, wavemax, expected",
+    [
+        pytest.param(
+            617,
+            615,
+            619,
+            "Fe I (617.33 nm)",
+            id="line_is_between_wavemin_and_wavemax_and_exists",
+        ),
+        pytest.param(
+            700,
+            698,
+            702,
+            None,
+            id="line_is_between_wavemin_and_wavemax_and_does_not_exist",
+        ),
+        pytest.param(
+            617,
+            698,
+            702,
+            None,
+            id="line_is_not_between_wavemin_and_wavemax_and_exists",
+        ),
+    ],
+)
+def test_get_waveband(write_l1_task, wavelength, wavemin, wavemax, expected):
+    """
+    :Given: an input wavelength contribution
+    :When: determining the waveband
+    :Then: the correct waveband is returned
+    """
+    wavelength_range = WavelengthRange(min=wavemin * u.nm, max=wavemax * u.nm)
+    task, _, _ = write_l1_task
+    waveband = task.get_waveband(wavelength=wavelength * u.nm, wavelength_range=wavelength_range)
+    assert waveband == expected
+
+
+def test_empty_waveband(write_l1_task_with_empty_waveband, mocker, fake_gql_client):
+    """
+    :Given: a header converted to SPEC 214 L1 and a wavelength range that has no listed spectral lines
+    :When: checking the waveband key
+    :Then: it does not exist
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
+    )
+    mocker.patch(
+        "dkist_processing_common.tasks.write_l1.WriteL1Frame.version_from_module_name",
+        new_callable=Mock,
+        return_value="fake_version_number",
+    )
+
+    task, _, _ = write_l1_task_with_empty_waveband
+    task()
+    files = list(task.read(tags=[Tag.frame(), Tag.output()]))
+    for file in files:
+        header = fits.open(file)[1].header
+        assert header["WAVEMIN"] == 10000
+        assert header["WAVEMAX"] == 10050
+        with pytest.raises(KeyError):
+            header["WAVEBAND"]