diff --git a/pyproject.toml b/pyproject.toml index 51defe6fe..f515682e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -119,6 +119,7 @@ TomographyMetadataContext = "murfey.client.contexts.tomo_metadata:TomographyMeta "experiment_type_update" = "murfey.workflows.register_experiment_type_update:run" "fib.make_milling_gif" = "murfey.workflows.fib.make_milling_gif:run" "fib.register_atlas" = "murfey.workflows.fib.register_atlas:run" +"fib.register_lamella_evaluation_image" = "murfey.workflows.fib.register_lamella_evaluation_image:run" "fib.register_milling_progress" = "murfey.workflows.fib.register_milling_progress:run" "pato" = "murfey.workflows.notifications:notification_setup" "picked_particles" = "murfey.workflows.spa.picking:particles_picked" diff --git a/src/murfey/client/contexts/fib.py b/src/murfey/client/contexts/fib.py index 48964ef21..8411a45f6 100644 --- a/src/murfey/client/contexts/fib.py +++ b/src/murfey/client/contexts/fib.py @@ -290,6 +290,15 @@ def post_transfer( ): self._make_drift_correction_gif(transferred_file, environment) return None + # Register lamella evaluation images + elif ( + "LamellaEvaluationImages" in transferred_file.parts + and transferred_file.suffix == ".png" + ): + self._register_lamella_evaluation_image( + transferred_file, environment + ) + return None # ----------------------------------------------------------------------------- # Maps @@ -300,19 +309,8 @@ def post_transfer( "Electron Snapshot" in transferred_file.name and transferred_file.suffix in (".tif", ".tiff") ): - source = _get_source(transferred_file, environment) - if source is None: - logger.warning(f"No source found for file {transferred_file}") - return None - destination_file = _file_transferred_to( - environment=environment, - source=source, - file_path=transferred_file, - rsync_basepath=Path(self._machine_config.get("rsync_basepath", "")), - ) - # Register image in database - self._register_atlas(destination_file, environment) + self._register_atlas(transferred_file, environment) return None # ----------------------------------------------------------------------------- @@ -651,12 +649,65 @@ def _make_gif( logger.error(f"Could not submit GIF for site {lamella_number}") return False - def _register_atlas(self, file: Path, environment: MurfeyInstanceEnvironment): + def _register_lamella_evaluation_image( + self, + transferred_file: Path, + environment: MurfeyInstanceEnvironment, + ): + """ + Helper function to construct and submit a POST request to the backend register + a lamella evaluation image with. + """ + source = _get_source(transferred_file, environment) + if source is None: + logger.warning(f"No source found for file {transferred_file}") + return None + destination_file = _file_transferred_to( + environment=environment, + source=source, + file_path=transferred_file, + rsync_basepath=Path(self._machine_config.get("rsync_basepath", "")), + ) + try: + capture_post( + base_url=str(environment.url.geturl()), + router_name="workflow_fib.router", + function_name="register_lamella_evaluation_image", + token=self._token, + instrument_name=environment.instrument_name, + data={"file": str(destination_file)}, + # Endpoint kwargs + session_id=environment.murfey_session, + ) + logger.info( + f"Registering lamella evaluation image {transferred_file.name!r}" + ) + return None + except Exception as e: + logger.error( + f"Error encountered registering lamella evaluation image {transferred_file.name}:\n{e}" + ) + return None + + def _register_atlas( + self, + transferred_file: Path, + environment: MurfeyInstanceEnvironment, + ): """ Constructs the URL and dictionary to be posted to the server, which then triggers the processing of the electron snapshot image. """ - + source = _get_source(transferred_file, environment) + if source is None: + logger.warning(f"No source found for file {transferred_file}") + return None + destination_file = _file_transferred_to( + environment=environment, + source=source, + file_path=transferred_file, + rsync_basepath=Path(self._machine_config.get("rsync_basepath", "")), + ) try: capture_post( base_url=str(environment.url.geturl()), @@ -664,12 +715,14 @@ def _register_atlas(self, file: Path, environment: MurfeyInstanceEnvironment): function_name="register_fib_atlas", token=self._token, instrument_name=environment.instrument_name, - data={"file": str(file)}, + data={"file": str(destination_file)}, # Endpoint kwargs session_id=environment.murfey_session, ) - logger.info(f"Registering atlas image {file.name!r}") - return True + logger.info(f"Registering atlas image {transferred_file.name!r}") + return None except Exception as e: - logger.error(f"Error encountered registering atlas image {file.name}:\n{e}") - return False + logger.error( + f"Error encountered registering atlas image {transferred_file.name}:\n{e}" + ) + return None diff --git a/src/murfey/server/api/workflow_fib.py b/src/murfey/server/api/workflow_fib.py index 1b02e6e7a..817d2203e 100644 --- a/src/murfey/server/api/workflow_fib.py +++ b/src/murfey/server/api/workflow_fib.py @@ -17,14 +17,14 @@ ) -class FIBAtlasFile(BaseModel): +class File(BaseModel): file: Path @router.post("/sessions/{session_id}/register_atlas") def register_fib_atlas( session_id: int, - fib_atlas: FIBAtlasFile, + fib_atlas: File, ): if _transport_object is None: logger.error("No TransportManager object was set up") @@ -39,6 +39,24 @@ def register_fib_atlas( ) +@router.post("/sessions/{session_id}/register_lamella_evaluation_image") +def register_lamella_evaluation_image( + session_id: int, + lamella_image: File, +): + if _transport_object is None: + logger.error("No TransportManager object was set up") + return None + _transport_object.send( + _transport_object.feedback_queue, + { + "register": "fib.register_lamella_evaluation_image", + "session_id": session_id, + "lamella_image_file": str(lamella_image.file), + }, + ) + + @router.post("/sessions/{session_id}/register_milling_progress") def register_fib_milling_progress( session_id: int, diff --git a/src/murfey/util/route_manifest.yaml b/src/murfey/util/route_manifest.yaml index bb8641d9e..0e5799fff 100644 --- a/src/murfey/util/route_manifest.yaml +++ b/src/murfey/util/route_manifest.yaml @@ -1431,6 +1431,13 @@ murfey.server.api.workflow_fib.router: type: int methods: - POST + - path: /workflow/fib/sessions/{session_id}/register_lamella_evaluation_image + function: register_lamella_evaluation_image + path_params: + - name: session_id + type: int + methods: + - POST - path: /workflow/fib/sessions/{session_id}/register_milling_progress function: register_fib_milling_progress path_params: diff --git a/src/murfey/workflows/fib/register_lamella_evaluation_image.py b/src/murfey/workflows/fib/register_lamella_evaluation_image.py new file mode 100644 index 000000000..2e826f16a --- /dev/null +++ b/src/murfey/workflows/fib/register_lamella_evaluation_image.py @@ -0,0 +1,17 @@ +import json +import logging +from typing import Any + +from sqlmodel import Session + +logger = logging.getLogger(__name__) + + +def run( + message: dict[str, Any], + murfey_db: Session, +): + logger.debug( + "Received message containing the following:\n" + f"{json.dumps(message, indent=2, default=str)}" + ) diff --git a/tests/client/contexts/test_fib.py b/tests/client/contexts/test_fib.py index 756362073..3fef13238 100644 --- a/tests/client/contexts/test_fib.py +++ b/tests/client/contexts/test_fib.py @@ -1,5 +1,6 @@ import xml.etree.ElementTree as ET from pathlib import Path +from typing import Any from unittest import mock from unittest.mock import MagicMock @@ -24,6 +25,8 @@ num_lamellae = 5 visit_name = "cm12345-6" project_name = visit_name.replace("-", "_") +session_id = 1 +instrument_name = "fib" # ------------------------------------------------------------------------------------- @@ -61,7 +64,7 @@ def visit_dir(tmp_path: Path): @pytest.fixture -def mock_machine_config(): +def mock_machine_config() -> dict[str, Any]: return {"calibrations": {"rotation_offset": -75}} @@ -849,12 +852,110 @@ def test_make_drift_correction_gif( assert mock_capture_post.call_count == len(destination_files) +@pytest.mark.parametrize( + "test_params", + ( # Has source | Is manual? | Site num | File part + # Successful cases + (True, True, 1, "Finer Milling - Electron Image.png"), + (True, False, 1, "Polishing 1 - Electron Image.png"), + (True, True, 2, "Finer Milling - Electron Image.png"), + (True, False, 2, "Polishing 1 - Electron Image.png"), + (True, True, 11, "Finer Milling - Electron Image.png"), + (True, False, 11, "Polishing 1 - Electron Image.png"), + # Early exit cases + (False, False, 1, "Finer Milling - Electron Image.png"), + (False, True, 1, "Polishing 1 - Electron Image.png"), + ), +) +def test_register_lamella_evaluation_image( + mocker: MockerFixture, + test_params: tuple[bool, bool, int, str], + tmp_path: Path, + visit_dir: Path, + mock_machine_config: dict[str, Any], +): + # Unpack test params + has_source, is_manual, site_num, file_part = test_params + + # Construct source and destination file names + if is_manual: + site_name = f"Site #{site_num}" + else: + site_name = "Lamella" + if site_num > 1: + site_name += f" ({site_num})" + transferred_file = ( + visit_dir + / "autotem" + / (f"AutoTEM_201231-1230_{project_name}_waffle1" if is_manual else project_name) + / "Sites" + / site_name + / "LamellaEvaluationImages" + / f"2025-10-23-19-14-40_drift_corrected_image_{file_part}" + ) + rsync_basepath = tmp_path / "fib" / "data" + destination_file = ( + rsync_basepath / "2025" / transferred_file.relative_to(visit_dir.parent) + ) + + # Mock the environment + mock_environment = MagicMock( + instrument_name=instrument_name, + murfey_session=session_id, + ) + # Mock '_get_source' + mock_get_source = mocker.patch( + "murfey.client.contexts.fib._get_source", + return_value=visit_dir if has_source else None, + ) + # Mock '_file_transferred_to' + mock_file_transferred_to = mocker.patch( + "murfey.client.contexts.fib._file_transferred_to", return_value=destination_file + ) + # Mock the 'capture_post' call + mock_capture_post = mocker.patch("murfey.client.contexts.fib.capture_post") + + # Load the context and pass in the file + basepath = tmp_path + context = FIBContext( + "autotem", + basepath=basepath, + machine_config=mock_machine_config, + token="dummy", + ) + context.post_transfer(transferred_file, mock_environment) + + # Check that the expected calls were made + mock_get_source.assert_called_once_with(transferred_file, mock_environment) + if not has_source: + mock_file_transferred_to.assert_not_called() + mock_capture_post.assert_not_called() + else: + mock_file_transferred_to.assert_called_once_with( + environment=mock_environment, + source=visit_dir, + file_path=transferred_file, + rsync_basepath=mock.ANY, + ) + mock_capture_post.assert_called_once_with( + base_url=mock.ANY, + router_name="workflow_fib.router", + function_name="register_lamella_evaluation_image", + token=context._token, + instrument_name=instrument_name, + data={"file": str(destination_file)}, + # Endpoint kwargs + session_id=session_id, + ) + + @pytest.mark.parametrize( "test_params", ( # Manual or automated? | Identifier (True, "Sites/Site #1/DCImages/dummy.png"), (True, "Sites/Site #1/LamellaEvaluationImages/dummy.png"), (False, "Sites/Lamella/DCImages/dummy.png"), + (False, "Sites/Lamella/LamellaEvaluationImages/dummy.png"), ), ) def test_fib_autotem_context( @@ -897,6 +998,9 @@ def test_fib_autotem_context( mock_drift_correction_gif = mocker.patch.object( FIBContext, "_make_drift_correction_gif" ) + mock_lamella_evaluation_image = mocker.patch.object( + FIBContext, "_register_lamella_evaluation_image" + ) # Initialise the FIBContext basepath = visit_dir @@ -921,6 +1025,9 @@ def test_fib_autotem_context( # If a DCImage was used, '_make_drift_correction_gif' should be called if "DCImages" in trigger_file.parts: mock_drift_correction_gif.assert_called_with(trigger_file, mock_environment) + # If a LamellaEvaluationImage was used '_register_lamella_evaluation_image' should be called + if "LamellaEvaluationImages" in trigger_file.parts: + mock_lamella_evaluation_image.assert_called_with(trigger_file, mock_environment) # Target project will have been identified assert _get_project_name(target_project) in context._target_projects # '_handle_metadata' will have been called @@ -933,8 +1040,13 @@ def test_fib_autotem_context( mock_handle_metadata.assert_called_with(target_project, mock_environment) +@pytest.mark.parametrize( + "has_source", + (True, False), +) def test_fib_maps_context( mocker: MockerFixture, + has_source: bool, tmp_path: Path, visit_dir: Path, mock_machine_config: dict, @@ -951,10 +1063,12 @@ def test_fib_maps_context( # Mock the functions used in 'post_transfer' mock_get_source = mocker.patch( - "murfey.client.contexts.fib._get_source", return_value=tmp_path + "murfey.client.contexts.fib._get_source", + return_value=tmp_path if has_source else None, ) mock_file_transferred_to = mocker.patch( - "murfey.client.contexts.fib._file_transferred_to", side_effect=destination_files + "murfey.client.contexts.fib._file_transferred_to", + side_effect=destination_files, ) mock_capture_post = mocker.patch("murfey.client.contexts.fib.capture_post") @@ -971,21 +1085,25 @@ def test_fib_maps_context( for f, file in enumerate(fib_maps_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(""), - ) - mock_capture_post.assert_called_with( - base_url=mock.ANY, - router_name="workflow_fib.router", - function_name="register_fib_atlas", - token="", - instrument_name=mock.ANY, - data={"file": str(destination_files[f])}, - session_id=mock.ANY, - ) + if has_source: + mock_file_transferred_to.assert_any_call( + environment=mock_environment, + source=basepath, + file_path=file, + rsync_basepath=Path(""), + ) + mock_capture_post.assert_any_call( + base_url=mock.ANY, + router_name="workflow_fib.router", + function_name="register_fib_atlas", + token="", + instrument_name=mock.ANY, + data={"file": str(destination_files[f])}, + session_id=mock.ANY, + ) + else: + mock_file_transferred_to.assert_not_called() + mock_capture_post.assert_not_called() def test_fib_meteor_context(): diff --git a/tests/server/api/test_workflow_fib.py b/tests/server/api/test_workflow_fib.py index 3a8f997b0..8d835283f 100644 --- a/tests/server/api/test_workflow_fib.py +++ b/tests/server/api/test_workflow_fib.py @@ -5,9 +5,10 @@ from pytest_mock import MockerFixture from murfey.server.api.workflow_fib import ( - FIBAtlasFile, + File, make_gif, register_fib_atlas, + register_lamella_evaluation_image, ) from murfey.util.models import FIBGIFParameters @@ -26,7 +27,7 @@ def test_register_fib_atlas( ): # Set up the variables session_id = 1 - fib_atlas = FIBAtlasFile(**{"file": str(tmp_path / "dummy")}) + fib_atlas = File(**{"file": str(tmp_path / "dummy")}) # Mock the logger mock_logger = mocker.patch("murfey.server.api.workflow_fib.logger") @@ -65,6 +66,59 @@ def test_register_fib_atlas( mock_logger.error.assert_called_with("No TransportManager object was set up") +@pytest.mark.parametrize( + "has_transport_object", + ( + True, + False, + ), +) +def test_register_lamella_evaluation_image( + mocker: MockerFixture, + tmp_path: Path, + has_transport_object: bool, +): + # Set up the variables + session_id = 1 + lamella_image = File(**{"file": str(tmp_path / "dummy")}) + + # Mock the logger + mock_logger = mocker.patch("murfey.server.api.workflow_fib.logger") + + # Mock the transport object + if has_transport_object: + mock_transport_object = MagicMock() + mock_transport_object.feedback_queue = "dummy" + mocker.patch( + "murfey.server.api.workflow_fib._transport_object", + mock_transport_object, + ) + else: + mocker.patch( + "murfey.server.api.workflow_fib._transport_object", + None, + ) + + # Run the function and check that the expected calls were made + register_lamella_evaluation_image( + session_id=session_id, + lamella_image=lamella_image, + ) + + # Check that the expected calls were made + if has_transport_object: + mock_transport_object.send.assert_called_with( + "dummy", + { + "register": "fib.register_lamella_evaluation_image", + "session_id": session_id, + "lamella_image_file": str(lamella_image.file), + }, + ) + else: + mock_logger.error.assert_called_with("No TransportManager object was set up") + + @pytest.mark.parametrize( "has_transport_object", ( diff --git a/tests/workflows/fib/test_register_lamella_evaluation_image.py b/tests/workflows/fib/test_register_lamella_evaluation_image.py new file mode 100644 index 000000000..07fbf67a7 --- /dev/null +++ b/tests/workflows/fib/test_register_lamella_evaluation_image.py @@ -0,0 +1,17 @@ +from unittest.mock import MagicMock + +from pytest_mock import MockerFixture + +from murfey.workflows.fib.register_lamella_evaluation_image import run + + +def test_run( + mocker: MockerFixture, +): + mock_logger = mocker.patch( + "murfey.workflows.fib.register_lamella_evaluation_image.logger" + ) + mock_murfey_db = MagicMock() + message = {"dummy": "dummy"} + run(message, mock_murfey_db) + mock_logger.debug.assert_called_once()