IMAT ngem transforms

.gitignore

@@ -1,4 +1,5 @@
 .idea/
+.junie/
 fia_api.egg-info/
 
 fia_api/local_scripts/*

fia_api/scripts/transforms/imat_transforms.py

@@ -4,31 +4,66 @@
 """
 
 import logging
+from pathlib import Path
 
+from fia_api.core.auth.tokens import DEV_MODE
 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__)
 
 
+def _extract_cycle_details(ngem_run_path: str) -> tuple[str, str]:
+    """An example path looks like: /ngem/nGEM-INES/DATA/IMAT_2026_01/IMAT00038896"""
+    ngem_run_parent_directory = Path(ngem_run_path).parent.name
+    _, cycle_year, cycle_num = ngem_run_parent_directory.split("_")
+    return cycle_num, cycle_year[-2:]
+
+
 class IMATTransform(Transform):
     """
     IMATTransform applies modifications to IMAT instrument scripts based on reduction input parameters in a Reduction
     entity.
     """
 
-    def apply(self, script: PreScript, job: Job) -> None:
+    def apply(self, script: PreScript, job: Job) -> None:  # noqa: C901, PLR0912
         logger.info("Beginning IMAT 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.
         for index, line in enumerate(lines):
-            if line.startswith("runno ="):
+            if line.startswith("runno =") and "runno" in job.inputs:  # type: ignore
                 lines[index] = f"runno = {job.inputs['runno']}"  # type: ignore
                 continue
-            if line.startswith("dataset_path = "):
-                lines[index] = f'dataset_path = Path("{job.inputs["images_dir"]}")'  # type: ignore
+            if line.startswith("dataset_path =") and "images_dir" in job.inputs:  # type: ignore
+                lines[index] = f'dataset_path = "{job.inputs["images_dir"]}"'  # type: ignore
+                continue
+            if line.startswith("ngem_path =") and "ngem_path" in job.inputs:  # type: ignore
+                lines[index] = f'ngem_path = "{job.inputs["ngem_path"]}"'  # type: ignore
+                continue
+            if line.startswith("ngem ="):
+                # Regardless we want to set the boolean state
+                if "ngem" in job.inputs and job.inputs["ngem"] == "true":  # type: ignore
+                    lines[index] = "ngem = True"
+                else:
+                    lines[index] = "ngem = False"
+                continue
+            if line.startswith("recon ="):
+                # Regardless we want to set the boolean state
+                if "recon" in job.inputs and job.inputs["recon"] == "true":  # type: ignore
+                    lines[index] = "recon = True"
+                else:
+                    lines[index] = "recon = False"
+                continue
+            if line.startswith("output ="):
+                if "ngem_path" in job.inputs and not DEV_MODE:  # type: ignore
+                    cycle_num, cycle_year = _extract_cycle_details(job.inputs["ngem_path"])  # type: ignore
+                    imat_nxs_folder = f"IMAT_{cycle_year}_{cycle_num}_nxs"
+                    output_path = f'"{Path(job.inputs["ngem_path"]).parent.parent / imat_nxs_folder}"'  # type: ignore
+                else:
+                    output_path = '"/output"'
+                lines[index] = f"output = {output_path}"
                 continue
 
         script.value = "\n".join(lines)

test/scripts/transforms/test_imat_transform.py

@@ -1,9 +1,9 @@
-"""Test Case for osiris transforms"""
+"""Test Case for IMAT transforms"""
 
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 from fia_api.scripts.pre_script import PreScript
-from fia_api.scripts.transforms.imat_transforms import IMATTransform
+from fia_api.scripts.transforms.imat_transforms import IMATTransform, _extract_cycle_details
 
 SCRIPT = """
 FILTERS = {f.__name__: f for f in load_filter_packages()}
@@ -15,27 +15,127 @@
 # To be edited by us for the script
 runno = 112345
 dataset_path = Path("/home/ubuntu/large")
+ngem_path = "/path/to/ngem"
+ngem = False
+recon = False
+output = "/output"
 """
 
 
-def test_iris_transform_spectroscopy():
-    """Test spectroscopy transform"""
-    expected_script = """
-FILTERS = {f.__name__: f for f in load_filter_packages()}
-RECON_DEFAULT_SETTINGS = {'algorithm': 'FBP_CUDA', 'filter_name': 'ram-lak', 'cor': 1, 'tilt': 0}
-DEBUG = False
-DEBUG_DIR = Path("/output/debug")
+def apply_transform_and_verify(job_inputs, expected_replacements, dev_mode=False):
+    """Helper to apply transform and verify results against expected replacements in SCRIPT"""
+    job = Mock()
+    job.id = "job_test"
+    job.inputs = job_inputs
+    script = PreScript(value=SCRIPT.strip())
+    with patch("fia_api.scripts.transforms.imat_transforms.DEV_MODE", dev_mode):
+        IMATTransform().apply(script, job)
 
+    expected_lines = SCRIPT.strip().splitlines()
 
-# To be edited by us for the script
-runno = 99999
-dataset_path = Path("/super/cool/path")"""
-    job = Mock()
-    job.inputs = {
+    for old, new in expected_replacements.items():
+        found = False
+        for i, line in enumerate(expected_lines):
+            if line.startswith(old):
+                expected_lines[i] = new
+                found = True
+        if not found:
+            # If not in SCRIPT but expected, it won't be in script.value either unless we add it
+            pass
+
+    assert script.value == "\n".join(expected_lines)
+
+
+def test_imat_transform_all_inputs():
+    """Test IMAT transform with all inputs provided"""
+    inputs = {
         "runno": 99999,
         "images_dir": "/super/cool/path",
+        "ngem_path": "/cycle/IMAT_2024_1/file.txt",
+        "ngem": "true",
+        "recon": "true",
+    }
+    replacements = {
+        "runno =": "runno = 99999",
+        "dataset_path =": 'dataset_path = "/super/cool/path"',
+        "ngem_path =": 'ngem_path = "/cycle/IMAT_2024_1/file.txt"',
+        "ngem =": "ngem = True",
+        "recon =": "recon = True",
+        "output =": 'output = "/cycle/IMAT_24_1_nxs"',
+    }
+    apply_transform_and_verify(inputs, replacements, dev_mode=False)
+
+
+def test_imat_transform_dev_mode():
+    """Test IMAT transform when DEV_MODE is True"""
+    inputs = {
+        "ngem_path": "/new/ngem/path/file.txt",
     }
-    script = PreScript(value=SCRIPT)
-    IMATTransform().apply(script, job)
+    replacements = {
+        "ngem_path =": 'ngem_path = "/new/ngem/path/file.txt"',
+        "ngem =": "ngem = False",
+        "recon =": "recon = False",
+        "output =": 'output = "/output"',
+    }
+    apply_transform_and_verify(inputs, replacements, dev_mode=True)
+
+
+def test_imat_transform_defaults():
+    """Test IMAT transform with default/false values for booleans and missing optional inputs"""
+    inputs = {
+        "ngem": "false",
+        "recon": "false",
+    }
+    replacements = {
+        "ngem =": "ngem = False",
+        "recon =": "recon = False",
+        "output =": 'output = "/output"',
+    }
+    apply_transform_and_verify(inputs, replacements)
+
+
+def test_imat_transform_missing_booleans():
+    """Test IMAT transform when boolean inputs are missing, they should default to False in script"""
+    inputs = {}
+    replacements = {
+        "ngem =": "ngem = False",
+        "recon =": "recon = False",
+        "output =": 'output = "/output"',
+    }
+    apply_transform_and_verify(inputs, replacements)
+
+
+def test_imat_transform_partial_inputs():
+    """Test IMAT transform with some inputs provided and some missing"""
+    inputs = {
+        "runno": 12345,
+        "ngem": "true",
+    }
+    replacements = {
+        "runno =": "runno = 12345",
+        "ngem =": "ngem = True",
+        "recon =": "recon = False",
+        "output =": 'output = "/output"',
+    }
+    apply_transform_and_verify(inputs, replacements)
+
+
+def test_extract_cycle_details_standard():
+    """Test _extract_cycle_details with a standard path format"""
+    cycle_num, cycle_year = _extract_cycle_details("/some/path/IMAT_2024_1/file.txt")
+    assert cycle_num == "1"
+    assert cycle_year == "24"
+
+
+def test_extract_cycle_details_different_year_and_cycle():
+    """Test _extract_cycle_details with a different year and cycle format"""
+    cycle_num, cycle_year = _extract_cycle_details("/data/IMAT_2023_02/data.nxs")
+    assert cycle_num == "02"
+    assert cycle_year == "23"
+
 
-    assert script.value == expected_script
+def test_extract_cycle_details_multiple_underscores():
+    """Test _extract_cycle_details with multiple underscores in the path"""
+    cycle_num, cycle_year = _extract_cycle_details("/path/to/some_folder/IMAT_2022_3/my_file_name.txt")
+    assert cycle_num == "3"
+    assert cycle_year == "22"