Added unit test for the FIBContext for the AutoTEM post transfer workflow

Author: tieneupin

tests/client/contexts/test_fib.py

@@ -14,20 +14,51 @@
 # -------------------------------------------------------------------------------------
 # FIBContext test utilty functions and fixtures
 # -------------------------------------------------------------------------------------
+num_lamellae = 5
 
 
 @pytest.fixture
 def visit_dir(tmp_path: Path):
     return tmp_path / "visit"
 
 
+@pytest.fixture
+def fib_autotem_dc_images(visit_dir: Path):
+    stages = ("Rough-Milling", "Polishing-1", "Polishing-2")
+    images_per_stage = 2
+
+    # Create images as Murfey would expect to find them in the DCImages folder
+    image_list = []
+    for l in range(num_lamellae):
+        for s, stage in enumerate(stages):
+            for i in range(images_per_stage):
+                lamella_folder = "Lamella"
+                if l > 0:
+                    lamella_folder += f" ({l + 1})"
+                # Continuously increment seconds count between files
+                timestamp = (
+                    f"2025-05-10-12-34-{str(0 + i + (s * images_per_stage)).zfill(2)}"
+                )
+                file = (
+                    visit_dir
+                    / f"autotem/visit/Sites/{lamella_folder}"
+                    / "DCImages/DCM_2025-05-10-12-34-00.125"
+                    / f"{timestamp}-{stage}-dc_rescan-image-.png"
+                )
+                if not file.exists():
+                    file.parent.mkdir(parents=True, exist_ok=True)
+                    file.touch()
+                image_list.append(file)
+    return image_list
+
+
 @pytest.fixture
 def fib_maps_images(visit_dir: Path):
     image_list = []
     for i in range(4):
         name = "Electron Snapshot"
         if i > 0:
-            name += f" ({i})"
+            name += f" ({i + 1})"
         file = visit_dir / "maps/visit/LayersData/Layer" / f"{name}.tiff"
         if not file.exists():
             file.parent.mkdir(parents=True, exist_ok=True)
@@ -97,8 +128,52 @@ def test_file_transferred_to(
         ) == destination_dir / file.relative_to(visit_dir)
 
 
-def test_fib_autotem_context():
-    pass
+def test_fib_autotem_context(
+    mocker: MockerFixture,
+    tmp_path: Path,
+    visit_dir: Path,
+    fib_autotem_dc_images: list[Path],
+):
+    # Mock the environment
+    mock_environment = MagicMock()
+
+    # Create a list of destinations
+    destination_dir = tmp_path / "fib" / "data" / "current_year" / "visit"
+    destination_files = [
+        destination_dir / file.relative_to(visit_dir) for file in fib_autotem_dc_images
+    ]
+
+    # Mock the functions used in 'post_transfer'
+    mock_get_source = mocker.patch(
+        "murfey.client.contexts.fib._get_source", return_value=tmp_path
+    )
+    mock_file_transferred_to = mocker.patch(
+        "murfey.client.contexts.fib._file_transferred_to", side_effect=destination_files
+    )
+    mock_capture_post = mocker.patch("murfey.client.contexts.fib.capture_post")
+
+    # Initialise the FIBContext
+    basepath = tmp_path
+    context = FIBContext(
+        acquisition_software="autotem",
+        basepath=basepath,
+        machine_config={},
+        token="",
+    )
+
+    # 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)
+        mock_get_source.assert_called_with(file, mock_environment)
+        mock_file_transferred_to.assert_called_with(
+            environment=mock_environment,
+            source=basepath,
+            file_path=file,
+            rsync_basepath=Path(""),
+        )
+    assert mock_capture_post.call_count == len(fib_autotem_dc_images)
+    assert len(context._milling) == num_lamellae
+    assert len(context._lamellae) == num_lamellae
 
 
 def test_fib_maps_context(
@@ -139,8 +214,13 @@ def test_fib_maps_context(
     # Parse images one-by-one
     for file in fib_maps_images:
         context.post_transfer(file, environment=mock_environment)
-    assert mock_get_source.call_count == len(fib_maps_images)
-    assert mock_file_transferred_to.call_count == len(fib_maps_images)
+        mock_get_source.assert_called_with(file, mock_environment)
+        mock_file_transferred_to.assert_called_with(
+            environment=mock_environment,
+            source=basepath,
+            file_path=file,
+            rsync_basepath=Path(""),
+        )
     assert mock_register_fib_atlas.call_count == len(fib_maps_images)
     for dst in destination_files:
         mock_register_fib_atlas.assert_any_call(