diff --git a/fia_api/scripts/transforms/factory.py b/fia_api/scripts/transforms/factory.py index 7dadf22b..05573597 100644 --- a/fia_api/scripts/transforms/factory.py +++ b/fia_api/scripts/transforms/factory.py @@ -3,6 +3,7 @@ import logging from fia_api.scripts.transforms.enginx_transform import EnginxTransform +from fia_api.scripts.transforms.gem_transform import GEMTransform from fia_api.scripts.transforms.imat_transforms import IMATTransform from fia_api.scripts.transforms.iris_transform import IrisTransform from fia_api.scripts.transforms.mari_transforms import MariTransform @@ -16,7 +17,7 @@ logger = logging.getLogger(__name__) -def get_transform_for_instrument(instrument: str) -> Transform: # noqa: PLR0911 +def get_transform_for_instrument(instrument: str) -> Transform: # noqa: PLR0911, C901 """ Get the appropriate transform for the given instrument and run file :param instrument: str - the instrument @@ -40,6 +41,8 @@ def get_transform_for_instrument(instrument: str) -> Transform: # noqa: PLR0911 return EnginxTransform() case "imat": return IMATTransform() + case "gem": + return GEMTransform() case "test": return TestTransform() case _: diff --git a/fia_api/scripts/transforms/gem_transform.py b/fia_api/scripts/transforms/gem_transform.py new file mode 100644 index 00000000..3536a813 --- /dev/null +++ b/fia_api/scripts/transforms/gem_transform.py @@ -0,0 +1,62 @@ +import logging + +from fia_api.core.models import Job +from fia_api.scripts.pre_script import PreScript +from fia_api.scripts.transforms.transform import Transform + +logger = logging.getLogger(__name__) + + +class GEMTransform(Transform): + """ + GEMTransform applies modifications to GEM instrument scripts based on reduction input parameters in a Reduction + entity. + """ + + def apply(self, script: PreScript, job: Job) -> None: # noqa: PLR0912,C901 + logger.info("Beginning GEM transform for job %s...", job.id) + lines = script.value.splitlines() + # MyPY does not believe ColumnElement[JSONB] is indexable, despite JSONB implementing the Indexable mixin + # If you get here in the future, try removing the following line and see if it passes with newer mypy. + + runno = job.inputs["runno"] # type: ignore + if isinstance(runno, list): + if len(runno) > 1: + # Convert list to range string if contiguous, otherwise comma-separated + if all(runno[i] == runno[i - 1] + 1 for i in range(1, len(runno))): + runno_str = f"{runno[0]}-{runno[-1]}" + else: + runno_str = ",".join(map(str, runno)) + else: + runno_str = str(runno[0]) + else: + runno_str = str(runno) + + for index, line in enumerate(lines): + if line.startswith("mode ="): + lines[index] = f'mode = "{job.inputs["mode"]}"' # type: ignore + continue + if line.startswith("input_mode ="): + lines[index] = f'input_mode = "{job.inputs["input_mode"]}"' # type: ignore + continue + if line.startswith("vanadium_runno ="): + lines[index] = f"vanadium_runno = {runno_str}" + continue + if line.startswith("runno ="): + lines[index] = f"runno = {runno_str}" + continue + if line.startswith("calibration_dir ="): + lines[index] = f"calibration_dir = {job.inputs['calibration_dir']}" # type: ignore + continue + if line.startswith("splined_vanadium_dir ="): + lines[index] = f'splined_vanadium_dir = "{job.inputs["splined_vanadium_dir"]}"' # type: ignore + continue + if line.startswith("config_file ="): + lines[index] = f'config_file = "{job.inputs["config_file"]}"' # type: ignore + continue + if line.startswith("output_dir = "): + lines[index] = f'output_dir = "{job.inputs["output_dir"]}"' # type: ignore + continue + + script.value = "\n".join(lines) + logger.info("Transform complete for job %s", job.id) diff --git a/test/scripts/transforms/test_factory.py b/test/scripts/transforms/test_factory.py index 6f1d673e..10524dd9 100644 --- a/test/scripts/transforms/test_factory.py +++ b/test/scripts/transforms/test_factory.py @@ -4,6 +4,7 @@ from fia_api.scripts.transforms.enginx_transform import EnginxTransform from fia_api.scripts.transforms.factory import get_transform_for_instrument +from fia_api.scripts.transforms.gem_transform import GEMTransform from fia_api.scripts.transforms.imat_transforms import IMATTransform from fia_api.scripts.transforms.iris_transform import IrisTransform from fia_api.scripts.transforms.mari_transforms import MariTransform @@ -28,6 +29,7 @@ ("vesuvio", VesuvioTransform), ("enginx", EnginxTransform), ("imat", IMATTransform), + ("gem", GEMTransform), ], ) def test_transform_factory(name, expected_transform): diff --git a/test/scripts/transforms/test_gem_transform.py b/test/scripts/transforms/test_gem_transform.py new file mode 100644 index 00000000..f1205f55 --- /dev/null +++ b/test/scripts/transforms/test_gem_transform.py @@ -0,0 +1,121 @@ +"""Test cases for GEMTransform.""" + +from unittest.mock import Mock + +import pytest + +from fia_api.scripts.pre_script import PreScript +from fia_api.scripts.transforms.gem_transform import GEMTransform + +SCRIPT = """ +mode = "default_mode" +input_mode = "default_input_mode" +vanadium_runno = 0 +runno = 0 +calibration_dir = None +splined_vanadium_dir = "default_splined_vanadium_dir" +config_file = "default_config_file" +output_dir = "default_output_dir" +""" + + +@pytest.fixture +def base_job(): + """Fixture for base job inputs.""" + job = Mock() + job.id = "test-job-gem" + job.inputs = { + "mode": "transmission", + "input_mode": "raw", + "calibration_dir": "/path/to/cal", + "splined_vanadium_dir": "/path/to/splined", + "config_file": "/path/to/config", + "output_dir": "/path/to/output", + "runno": 12345, + } + return job + + +@pytest.fixture +def create_expected_script(): + """Fixture returning a helper function to construct the expected script with runno.""" + + def _create(runno_str: str) -> str: + return f""" +mode = "transmission" +input_mode = "raw" +vanadium_runno = {runno_str} +runno = {runno_str} +calibration_dir = /path/to/cal +splined_vanadium_dir = "/path/to/splined" +config_file = "/path/to/config" +output_dir = "/path/to/output\"""" + + return _create + + +def test_gem_transform_single_run(base_job, create_expected_script): + """Test GEMTransform with a single run number.""" + script = PreScript(value=SCRIPT) + GEMTransform().apply(script, base_job) + + assert script.value == create_expected_script("12345") + + +def test_gem_transform_contiguous_runs(base_job, create_expected_script): + """Test GEMTransform with contiguous runs.""" + base_job.inputs["runno"] = [12345, 12346, 12347] + script = PreScript(value=SCRIPT) + GEMTransform().apply(script, base_job) + + assert script.value == create_expected_script("12345-12347") + + +def test_gem_transform_non_contiguous_runs(base_job, create_expected_script): + """Test GEMTransform with non-contiguous runs.""" + base_job.inputs["runno"] = [12345, 12347, 12349] + script = PreScript(value=SCRIPT) + GEMTransform().apply(script, base_job) + + assert script.value == create_expected_script("12345,12347,12349") + + +def test_gem_transform_list_length_one(base_job, create_expected_script): + """Test GEMTransform with list containing a single run.""" + base_job.inputs["runno"] = [12345] + script = PreScript(value=SCRIPT) + GEMTransform().apply(script, base_job) + + assert script.value == create_expected_script("12345") + + +def test_gem_transform_apply(base_job): + """Test GEMTransform only modifies expected lines and leaves others unchanged.""" + transform = GEMTransform() + script = PreScript(value=SCRIPT) + original_lines = script.value.splitlines() + + transform.apply(script, base_job) + + updated_lines = script.value.splitlines() + assert len(original_lines) == len(updated_lines) + + for index, line in enumerate(updated_lines): + if line.startswith("mode ="): + assert line == 'mode = "transmission"' + elif line.startswith("input_mode ="): + assert line == 'input_mode = "raw"' + elif line.startswith("vanadium_runno ="): + assert line == "vanadium_runno = 12345" + elif line.startswith("runno ="): + assert line == "runno = 12345" + elif line.startswith("calibration_dir ="): + assert line == "calibration_dir = /path/to/cal" + elif line.startswith("splined_vanadium_dir ="): + assert line == 'splined_vanadium_dir = "/path/to/splined"' + elif line.startswith("config_file ="): + assert line == 'config_file = "/path/to/config"' + elif line.startswith("output_dir = "): + assert line == 'output_dir = "/path/to/output"' + else: + assert line == original_lines[index]