Updated test for FIBContext to handle new drift correction FIB image logic

Author: tieneupin

tests/client/contexts/test_fib.py

@@ -16,6 +16,7 @@
     _number_from_name,
     _parse_boolean,
 )
+from murfey.util.models import LamellaSiteInfo
 
 # Mock session values
 num_lamellae = 5
@@ -595,27 +596,49 @@ def test_fib_autotem_context_projectdata(
 
 @pytest.mark.parametrize(
     "test_params",
-    (  # Use environment? | Find source? | Find destination?
-        (True, True, True),
-        (False, True, True),
-        (True, False, True),
-        (True, True, False),
+    (
+        # Early exits
+        # No MurfeyInstanceEnvironment
+        (False, True, True, True, True, True, True),
+        # No source
+        (True, False, True, True, True, True, True),
+        # No destination
+        (True, True, False, True, True, True, True),
+        # No site info
+        (True, True, True, False, True, True, True),
+        # No project name
+        (True, True, True, True, False, True, True),
+        # No stage position
+        (True, True, True, True, True, False, True),
+        # No stage position values
+        (True, True, True, True, True, True, False),
+        # Successful case
+        (True, True, True, True, True, True, True),
     ),
 )
 def test_fib_autotem_context_drift_correction_images(
     mocker: MockerFixture,
-    test_params: tuple[bool, bool, bool],
+    test_params: tuple[bool, bool, bool, bool, bool, bool, bool],
     tmp_path: Path,
     visit_dir: Path,
     fib_autotem_dc_images: list[Path],
 ):
     # Unpack test params
-    use_env, find_source, find_dst = test_params
+    (
+        use_env,
+        find_source,
+        find_dst,
+        has_site_info,
+        has_project_name,
+        has_stage_position,
+        has_stage_values,
+    ) = test_params
 
     # Mock the environment
     mock_environment = None
     if use_env:
         mock_environment = MagicMock()
+        mock_environment.visit = visit_name
 
     # Mock the logger to check if specific logs are triggered
     mock_logger = mocker.patch("murfey.client.contexts.fib.logger")
@@ -649,6 +672,23 @@ def test_fib_autotem_context_drift_correction_images(
         token="",
     )
 
+    # Create the Pydantic model for each site and add metadata
+    for i in range(num_lamellae):
+        lamella_num = i + 1
+        metadata_dict = {
+            "site_name": f"Lamella ({lamella_num})",
+            "site_number": lamella_num,
+        }
+        if has_project_name:
+            metadata_dict["project_name"] = project_name
+        if has_stage_position:
+            stage_dict: dict[str, dict] = {"preparation": {}}
+            if has_stage_values:
+                stage_dict["preparation"] = {"x": 0.003}
+            metadata_dict["stage_info"] = stage_dict
+        if has_site_info:
+            context._site_info[lamella_num] = LamellaSiteInfo(**metadata_dict)
+
     # Parse images one-by-one and check that expected calls were made
     for file in fib_autotem_dc_images:
         context.post_transfer(file, environment=mock_environment)
@@ -660,6 +700,22 @@ def test_fib_autotem_context_drift_correction_images(
         mock_logger.warning.assert_called_with(
             f"File {file.name!r} not found on storage system"
         )
+    elif not has_site_info:
+        mock_logger.debug.assert_called_with(
+            f"No metadata found for site {lamella_num} yet"
+        )
+    elif not has_project_name:
+        mock_logger.warning.assert_called_with(
+            f"No project name associated with site {lamella_num}"
+        )
+    elif not has_stage_position:
+        mock_logger.warning.assert_called_with(
+            f"No stage position information associated with site {lamella_num}"
+        )
+    elif not has_stage_values:
+        mock_logger.warning.assert_called_with(
+            f"Could not determine slot number of site {lamella_num}"
+        )
     else:
         mock_get_source.assert_called_with(file, mock_environment)
         mock_file_transferred_to.assert_called_with(
@@ -668,9 +724,34 @@ def test_fib_autotem_context_drift_correction_images(
             file_path=file,
             rsync_basepath=Path(""),
         )
-        assert mock_capture_post.call_count == len(fib_autotem_dc_images)
         assert len(context._drift_correction_images) == num_lamellae
 
+        for i in range(num_lamellae):
+            lamella_num = i + 1
+            # The '_site_info' attribute should now be populated
+            assert (
+                context._site_info[lamella_num].stage_info.preparation.slot_number == 2
+            )
+
+            # The output file should point to 'grid_2' for a positive x stage position
+            output_file = (
+                tmp_path
+                / "fib"
+                / "data"
+                / "current_year"
+                / visit_name
+                / "processed"
+                / project_name
+                / "grid_2"
+                / "drift_correction"
+                / f"lamella_{lamella_num}.gif"
+            )
+            assert (
+                context._drift_correction_images[lamella_num].output_file == output_file
+            )
+        # 'capture_post' should be called for every image
+        assert mock_capture_post.call_count == len(destination_files)
+
 
 def test_fib_maps_context(
     mocker: MockerFixture,