From 759f1e7cc73c80a90cfe602cd45249231e47e0f8 Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Mon, 1 Jun 2026 10:57:34 +0100 Subject: [PATCH 01/19] Added new Murfey workflow stub for registering FIB milling progress that can be run through the 'feedback_callback' function --- pyproject.toml | 1 + src/murfey/server/api/workflow_fib.py | 19 ++++++++++------ .../fib/register_milling_progress.py | 22 +++++++++++++++++++ 3 files changed, 35 insertions(+), 7 deletions(-) create mode 100644 src/murfey/workflows/fib/register_milling_progress.py diff --git a/pyproject.toml b/pyproject.toml index b181e111b..53645f58c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -117,6 +117,7 @@ TomographyMetadataContext = "murfey.client.contexts.tomo_metadata:TomographyMeta "data_collection_group" = "murfey.workflows.register_data_collection_group:run" "experiment_type_update" = "murfey.workflows.register_experiment_type_update:run" "fib.register_atlas" = "murfey.workflows.fib.register_atlas:run" +"fib.register_milling_progress" = "murfey.workflows.register_milling_progress:run" "pato" = "murfey.workflows.notifications:notification_setup" "picked_particles" = "murfey.workflows.spa.picking:particles_picked" "picked_tomogram" = "murfey.workflows.tomo.picking:picked_tomogram" diff --git a/src/murfey/server/api/workflow_fib.py b/src/murfey/server/api/workflow_fib.py index 577a11a4d..5c143a1f5 100644 --- a/src/murfey/server/api/workflow_fib.py +++ b/src/murfey/server/api/workflow_fib.py @@ -1,4 +1,3 @@ -import json import logging import os from pathlib import Path @@ -7,7 +6,7 @@ import PIL.Image from fastapi import APIRouter, Depends from pydantic import BaseModel -from sqlmodel import Session, select +from sqlmodel import select import murfey.util.db as MurfeyDB from murfey.server import _transport_object @@ -36,7 +35,7 @@ def register_fib_atlas( fib_atlas: FIBAtlasFile, ): if _transport_object is None: - logger.error("No Transport Manager object was set up") + logger.error("No TransportManager object was set up") return None _transport_object.send( _transport_object.feedback_queue, @@ -52,11 +51,17 @@ def register_fib_atlas( def register_fib_milling_progress( session_id: int, site_info: LamellaSiteInfo, - db: Session = murfey_db, ): - logger.debug( - "Received the following FIB metadata for registration:\n" - f"{json.dumps(site_info.model_dump(exclude_none=True), indent=2, default=str)}" + 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_milling_progress", + "session_id": session_id, + "site_info": site_info.model_dump(exclude_none=True), + }, ) diff --git a/src/murfey/workflows/fib/register_milling_progress.py b/src/murfey/workflows/fib/register_milling_progress.py new file mode 100644 index 000000000..3e78e8144 --- /dev/null +++ b/src/murfey/workflows/fib/register_milling_progress.py @@ -0,0 +1,22 @@ +import json +import logging +from typing import Any + +from sqlmodel import Session as SQLModelSession + +from murfey.util.models import LamellaSiteInfo + +logger = logging.getLogger("murfey.workflows.fib.register_milling_progress") + + +def run(message: dict[str, Any], murfey_db: SQLModelSession): + try: + session_id = int(message["session_id"]) + site_info = LamellaSiteInfo(**message["site_info"]) + logger.debug( + "Received the following FIB metadata for registration:\n" + f"{json.dumps(site_info.model_dump(exclude_none=True), indent=2, default=str)}" + ) + except Exception: + logger.error("Error parsing contents of message", exc_info=True) + return {"success": False, "requeue": False} From 4f4e128e7097a7714c4370daa54510b97a826ac3 Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Mon, 1 Jun 2026 10:57:48 +0100 Subject: [PATCH 02/19] Updated log string to check for --- tests/server/api/test_workflow_fib.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/server/api/test_workflow_fib.py b/tests/server/api/test_workflow_fib.py index 28b37ae29..1fd400c4d 100644 --- a/tests/server/api/test_workflow_fib.py +++ b/tests/server/api/test_workflow_fib.py @@ -64,7 +64,7 @@ def test_register_fib_atlas( }, ) else: - mock_logger.error.assert_called_with("No Transport Manager object was set up") + mock_logger.error.assert_called_with("No TransportManager object was set up") @pytest.mark.asyncio From aa9fb1ae5196c47fd3cbdb48508e5789560a99eb Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Mon, 1 Jun 2026 11:23:59 +0100 Subject: [PATCH 03/19] Wrong module path in 'pyproject.toml' --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 53645f58c..ce7398044 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -117,7 +117,7 @@ TomographyMetadataContext = "murfey.client.contexts.tomo_metadata:TomographyMeta "data_collection_group" = "murfey.workflows.register_data_collection_group:run" "experiment_type_update" = "murfey.workflows.register_experiment_type_update:run" "fib.register_atlas" = "murfey.workflows.fib.register_atlas:run" -"fib.register_milling_progress" = "murfey.workflows.register_milling_progress: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" "picked_tomogram" = "murfey.workflows.tomo.picking:picked_tomogram" From d6489fbd02df4de9f20d10e00a683efc4e66990a Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Mon, 1 Jun 2026 13:14:22 +0100 Subject: [PATCH 04/19] Bumped required 'ispyb' version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index ce7398044..5e8e46951 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,7 +55,7 @@ developer = [ server = [ "cryptography", "graypy", - "ispyb>=10.2.4", # Responsible for setting requirements for SQLAlchemy and mysql-connector-python; + "ispyb>=12.1.0", # Responsible for setting requirements for SQLAlchemy and mysql-connector-python; "jinja2", "mrcfile", "numpy<3", From e976593c86ae28beef9beeee65b42131f8d7b562 Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Tue, 2 Jun 2026 15:10:05 +0100 Subject: [PATCH 05/19] Added new TransportManager functinos to insert and update MillingStep database entries --- src/murfey/server/ispyb.py | 97 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/src/murfey/server/ispyb.py b/src/murfey/server/ispyb.py index 815f06354..310688106 100644 --- a/src/murfey/server/ispyb.py +++ b/src/murfey/server/ispyb.py @@ -20,6 +20,7 @@ DataCollectionGroup, FoilHole, GridSquare, + MillingStep, ProcessingJob, ProcessingJobParameter, Proposal, @@ -710,6 +711,102 @@ def do_update_processing_status(self, record: AutoProcProgram, **kwargs): ) return {"success": False, "return_value": None} + def do_insert_milling_step(self, record: MillingStep): + try: + with ISPyBSession() as db: + db.add(record) + db.commit() + log.info(f"Created MillingStep {record.millingStepId}") + return {"success": True, "return_value": record.millingStepId} + except ispyb.ISPyBException as e: + log.error( + "Insert MillingStep entry caused exception '%s'.", + e, + exc_info=True, + ) + return {"success": False, "return_value": None} + + def do_update_milling_step( + self, + milling_step_id: int, + is_enabled: bool | None = None, + status: str | None = None, + execution_time: float | None = None, + stage_x: float | None = None, + stage_y: float | None = None, + stage_z: float | None = None, + rotation: float | None = None, + tilt_alpha: float | None = None, + beam_type: str | None = None, + beam_voltage: float | None = None, + beam_current: float | None = None, + milling_angle: float | None = None, + depth_correction: float | None = None, + lamella_offset: float | None = None, + trench_height_front: float | None = None, + trench_height_rear: float | None = None, + width_overlap_front_left: float | None = None, + width_overlap_front_right: float | None = None, + width_overlap_rear_left: float | None = None, + width_overlap_rear_right: float | None = None, + ): + try: + with ISPyBSession() as db: + milling_step = ( + db.query(MillingStep) + .filter(MillingStep.millingStepId == milling_step_id) + .one() + ) + milling_step.isEnabled = is_enabled or milling_step.isEnabled + milling_step.status = status or milling_step.status + milling_step.executionTime = ( + execution_time or milling_step.executionTime + ) + milling_step.stageX = stage_x or milling_step.stageX + milling_step.stageY = stage_y or milling_step.stageY + milling_step.stageZ = stage_z or milling_step.stageZ + milling_step.rotation = rotation or milling_step.rotation + milling_step.alphaTilt = tilt_alpha or milling_step.alphaTilt + milling_step.beamType = beam_type or milling_step.beamType + milling_step.beamVoltage = beam_voltage or milling_step.beamVoltage + milling_step.beamCurrent = beam_current or milling_step.beamCurrent + milling_step.millingAngle = milling_angle or milling_step.millingAngle + milling_step.depthCorrection = ( + depth_correction or milling_step.depthCorrection + ) + milling_step.lamellaOffset = ( + lamella_offset or milling_step.lamellaOffset + ) + milling_step.trenchHeightFront = ( + trench_height_front or milling_step.trenchHeightFront + ) + milling_step.trenchHeightRear = ( + trench_height_rear or milling_step.trenchHeightRear + ) + milling_step.widthOverlapFrontLeft = ( + width_overlap_front_left or milling_step.widthOverlapFrontLeft + ) + milling_step.widthOverlapFrontRight = ( + width_overlap_front_right or milling_step.widthOverlapFrontRight + ) + milling_step.widthOverlapRearLeft = ( + width_overlap_rear_left or milling_step.widthOverlapRearLeft + ) + milling_step.widthOverlapRearRight = ( + width_overlap_rear_right or milling_step.widthOverlapRearRight + ) + + db.add(milling_step) + db.commit() + return {"success": True, "return_value": milling_step.millingStepId} + except ispyb.ISPyBException as e: + log.error( + "Updating MillingStep entry caused exception '%s'.", + e, + exc_info=True, + ) + return {"success": False, "return_value": None} + def do_buffer_lookup(self, app_id: int, uuid: int) -> Optional[int]: with ISPyBSession() as db: buffer_objects = ( From ec1cb851320f5142fa3f7225d4015febfcc2db94 Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Tue, 2 Jun 2026 15:36:54 +0100 Subject: [PATCH 06/19] Moved 'DataCollectionGroup' and 'GridSquare' to the 'General' section, and added a new table to record a Murfey-side copy of ISPyB's 'MillingStep' table --- src/murfey/util/db.py | 214 ++++++++++++++++++++++++++---------------- 1 file changed, 133 insertions(+), 81 deletions(-) diff --git a/src/murfey/util/db.py b/src/murfey/util/db.py index eb92f48fa..413336694 100644 --- a/src/murfey/util/db.py +++ b/src/murfey/util/db.py @@ -174,6 +174,139 @@ class ImagingSite(SQLModel, table=True): # type: ignore grid_square_id: Optional[int] = Field(foreign_key="gridsquare.id", default=None) +class DataCollectionGroup(SQLModel, table=True): # type: ignore + id: int = Field( + primary_key=True, + unique=True, + alias="dataCollectionGroupId", + sa_column_kwargs={"name": "dataCollectionGroupId"}, + ) + session_id: int = Field(foreign_key="session.id", primary_key=True) + tag: str = Field(primary_key=True) + atlas_id: Optional[int] = None + atlas_pixel_size: Optional[float] = None + atlas: str = "" + sample: Optional[int] = None + smartem_grid_uuid: Optional[str] = None + + # ------------- + # Relationships + # ------------- + session: Optional["Session"] = Relationship(back_populates="data_collection_groups") + data_collections: List["DataCollection"] = Relationship( + back_populates="data_collection_group", + sa_relationship_kwargs={"cascade": "delete"}, + ) + imaging_sites: List["ImagingSite"] = Relationship( + back_populates="data_collection_group", + sa_relationship_kwargs={"cascade": "delete"}, + ) + notification_parameters: List["NotificationParameter"] = Relationship( + back_populates="data_collection_group", + sa_relationship_kwargs={"cascade": "delete"}, + ) + tomography_processing_parameters: List["TomographyProcessingParameters"] = ( + Relationship( + back_populates="data_collection_group", + sa_relationship_kwargs={"cascade": "delete"}, + ) + ) + grid_squares: Optional[List["GridSquare"]] = Relationship( + back_populates="data_collection_group", + sa_relationship_kwargs={"cascade": "delete"}, + ) + search_maps: Optional[List["SearchMap"]] = Relationship( + back_populates="data_collection_group", + sa_relationship_kwargs={"cascade": "delete"}, + ) + + +class GridSquare(SQLModel, table=True): # type: ignore + id: Optional[int] = Field(primary_key=True, default=None) + session_id: int = Field(foreign_key="session.id") + name: int + tag: str + x_location: Optional[float] + y_location: Optional[float] + x_stage_position: Optional[float] + y_stage_position: Optional[float] + readout_area_x: Optional[int] + readout_area_y: Optional[int] + thumbnail_size_x: Optional[int] + thumbnail_size_y: Optional[int] + pixel_size: Optional[float] = None + image: str = "" + atlas_id: Optional[int] = Field( + foreign_key="datacollectiongroup.dataCollectionGroupId" + ) + scaled_pixel_size: Optional[float] = None + pixel_location_x: Optional[int] = None + pixel_location_y: Optional[int] = None + height: Optional[int] = None + width: Optional[int] = None + angle: Optional[float] = None + quality_indicator: Optional[float] = None + smartem_uuid: Optional[str] = None + + # ------------- + # Relationships + # ------------- + session: Optional[Session] = Relationship(back_populates="grid_squares") + imaging_sites: List["ImagingSite"] = Relationship( + back_populates="grid_square", sa_relationship_kwargs={"cascade": "delete"} + ) + milling_steps: List["MillingStep"] = Relationship( + back_populates="grid_square", sa_relationship_kwargs={"cascade": "delete"} + ) + foil_holes: List["FoilHole"] = Relationship( + back_populates="grid_square", sa_relationship_kwargs={"cascade": "delete"} + ) + data_collection_group: Optional["DataCollectionGroup"] = Relationship( + back_populates="grid_squares" + ) + + +""" +======================================================================================= +FIB WORKFLOW +======================================================================================= +""" + + +class MillingStep(SQLModel, table=True): # type: ignore + id: int = Field(primary_key=True, default=None) + recipe_name: Optional[str] = Field(default=None) + activity_name: Optional[str] = Field(default=None) + is_enabled: Optional[bool] = Field(default=None) + status: Optional[str] = Field(default=None) + execution_time: Optional[float] = Field(default=None) + stage_x: Optional[float] = Field(default=None) + stage_y: Optional[float] = Field(default=None) + stage_z: Optional[float] = Field(default=None) + rotation: Optional[float] = Field(default=None) + tilt_alpha: Optional[float] = Field(default=None) + beam_type: Optional[str] = Field(default=None) + beam_voltage: Optional[float] = Field(default=None) + beam_current: Optional[float] = Field(default=None) + milling_angle: Optional[float] = Field(default=None) + depth_correction: Optional[float] = Field(default=None) + lamella_offset: Optional[float] = Field(default=None) + trench_height_front: Optional[float] = Field(default=None) + trench_height_rear: Optional[float] = Field(default=None) + width_overlap_front_left: Optional[float] = Field(default=None) + width_overlap_front_right: Optional[float] = Field(default=None) + width_overlap_rear_left: Optional[float] = Field(default=None) + width_overlap_rear_right: Optional[float] = Field(default=None) + + # ------------- + # Relationships + # ------------- + grid_square: Optional["GridSquare"] = Relationship( + back_populates="milling_steps" + ) # Many-to-one + grid_square_id: Optional[int] = Field(foreign_key="gridsquare.id", default=None) + + """ ======================================================================================= TEM SESSION AND PROCESSING WORKFLOW @@ -223,49 +356,6 @@ class Tilt(SQLModel, table=True): # type: ignore tilt_series: Optional[TiltSeries] = Relationship(back_populates="tilts") -class DataCollectionGroup(SQLModel, table=True): # type: ignore - id: int = Field( - primary_key=True, - unique=True, - alias="dataCollectionGroupId", - sa_column_kwargs={"name": "dataCollectionGroupId"}, - ) - session_id: int = Field(foreign_key="session.id", primary_key=True) - tag: str = Field(primary_key=True) - atlas_id: Optional[int] = None - atlas_pixel_size: Optional[float] = None - atlas: str = "" - sample: Optional[int] = None - smartem_grid_uuid: Optional[str] = None - session: Optional["Session"] = Relationship(back_populates="data_collection_groups") - data_collections: List["DataCollection"] = Relationship( - back_populates="data_collection_group", - sa_relationship_kwargs={"cascade": "delete"}, - ) - imaging_sites: List["ImagingSite"] = Relationship( - back_populates="data_collection_group", - sa_relationship_kwargs={"cascade": "delete"}, - ) - notification_parameters: List["NotificationParameter"] = Relationship( - back_populates="data_collection_group", - sa_relationship_kwargs={"cascade": "delete"}, - ) - tomography_processing_parameters: List["TomographyProcessingParameters"] = ( - Relationship( - back_populates="data_collection_group", - sa_relationship_kwargs={"cascade": "delete"}, - ) - ) - grid_squares: Optional[List["GridSquare"]] = Relationship( - back_populates="data_collection_group", - sa_relationship_kwargs={"cascade": "delete"}, - ) - search_maps: Optional[List["SearchMap"]] = Relationship( - back_populates="data_collection_group", - sa_relationship_kwargs={"cascade": "delete"}, - ) - - class NotificationParameter(SQLModel, table=True): # type: ignore id: Optional[int] = Field(default=None, primary_key=True) dcg_id: int = Field(foreign_key="datacollectiongroup.dataCollectionGroupId") @@ -482,44 +572,6 @@ class MurfeyLedger(SQLModel, table=True): # type: ignore ) -class GridSquare(SQLModel, table=True): # type: ignore - id: Optional[int] = Field(primary_key=True, default=None) - session_id: int = Field(foreign_key="session.id") - name: int - tag: str - x_location: Optional[float] - y_location: Optional[float] - x_stage_position: Optional[float] - y_stage_position: Optional[float] - readout_area_x: Optional[int] - readout_area_y: Optional[int] - thumbnail_size_x: Optional[int] - thumbnail_size_y: Optional[int] - pixel_size: Optional[float] = None - image: str = "" - session: Optional[Session] = Relationship(back_populates="grid_squares") - imaging_sites: List["ImagingSite"] = Relationship( - back_populates="grid_square", sa_relationship_kwargs={"cascade": "delete"} - ) - foil_holes: List["FoilHole"] = Relationship( - back_populates="grid_square", sa_relationship_kwargs={"cascade": "delete"} - ) - atlas_id: Optional[int] = Field( - foreign_key="datacollectiongroup.dataCollectionGroupId" - ) - scaled_pixel_size: Optional[float] = None - pixel_location_x: Optional[int] = None - pixel_location_y: Optional[int] = None - height: Optional[int] = None - width: Optional[int] = None - angle: Optional[float] = None - quality_indicator: Optional[float] = None - smartem_uuid: Optional[str] = None - data_collection_group: Optional["DataCollectionGroup"] = Relationship( - back_populates="grid_squares" - ) - - class FoilHole(SQLModel, table=True): # type: ignore id: Optional[int] = Field(primary_key=True, default=None) grid_square_id: int = Field(foreign_key="gridsquare.id") From fd1493eda079f73fe4881ae51309b5ed0855aead Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Tue, 2 Jun 2026 15:37:52 +0100 Subject: [PATCH 07/19] Added logic to register FIB milling metadata in ISPyB and Murfey --- .../fib/register_milling_progress.py | 413 +++++++++++++++++- 1 file changed, 403 insertions(+), 10 deletions(-) diff --git a/src/murfey/workflows/fib/register_milling_progress.py b/src/murfey/workflows/fib/register_milling_progress.py index 3e78e8144..c745676bc 100644 --- a/src/murfey/workflows/fib/register_milling_progress.py +++ b/src/murfey/workflows/fib/register_milling_progress.py @@ -1,22 +1,415 @@ import json import logging +from importlib.metadata import entry_points from typing import Any -from sqlmodel import Session as SQLModelSession +import ispyb.sqlalchemy._auto_db_schema as ISPyBDB +from sqlmodel import Session as SQLModelSession, select -from murfey.util.models import LamellaSiteInfo +import murfey.util.db as MurfeyDB +from murfey.server import _transport_object +from murfey.util.fib import number_from_name +from murfey.util.models import ( + GridSquareParameters, + LamellaSiteInfo, + MillingStepInfo, + StagePositionValues, +) logger = logging.getLogger("murfey.workflows.fib.register_milling_progress") +def _ensure_prerequisites( + session_id: int, + instrument_name: str, + visit_name: str, + site_info: LamellaSiteInfo, + murfey_db: SQLModelSession, +): + """ + Uses the FIB milling metadata provided to create the necessary DataCollectionGroup, + Atlas, and GridSquare placeholders in Murfey if they don't already exist + + """ + # Skip this step if no TransportManager object is configured + if _transport_object is None: + logger.error("No TransportManager object was configured") + return None + + # Early exits if information needed to construct lookup tags are missing + if site_info.project_name is None: + logger.error("Could not construct lookup tags; 'project_name' is missing") + return None + if site_info.stage_info is None: + logger.error("Could not construct lookup tags; 'stage_info' is missing") + return None + if site_info.stage_info.preparation_site is None: + logger.error("Could not construct lookup tags; 'preparation_site' is missing") + return None + if site_info.stage_info.preparation_site.slot_number is None: + logger.error("Could not construct lookup tags; 'slot_number' is missing") + return None + slot_number = site_info.stage_info.preparation_site.slot_number + if site_info.site_name is None: + logger.error("Could not construct lookup tags; 'site_name' is missing") + return None + + # Construct the DataCollectionGroup and GridSquare lookup tags + lamella_num = number_from_name(site_info.site_name) + dcg_tag = f"{site_info.project_name}--slot_{slot_number}" + + # Determine variables to register data collection group and atlas with + proposal_code = "".join(char for char in visit_name.split("-")[0] if char.isalpha()) + proposal_number = "".join( + char for char in visit_name.split("-")[0] if char.isdigit() + ) + visit_number = visit_name.split("-")[-1] + + # Register the DataCollectionGroup and Atlas placeholder if it doesn't already exist + if ( + murfey_db.exec( + select(MurfeyDB.DataCollectionGroup) + .where(MurfeyDB.DataCollectionGroup.session_id == session_id) + .where(MurfeyDB.DataCollectionGroup.tag == dcg_tag) + ).one_or_none() + is None + ): + dcg_message = { + "microscope": instrument_name, + "proposal_code": proposal_code, + "proposal_number": proposal_number, + "visit_number": visit_number, + "session_id": session_id, + "tag": dcg_tag, + "experiment_type_id": 46, + "atlas": "", + "atlas_pixel_size": 0.0, + "sample": slot_number, + "color_flags": None, + "collection_mode": None, + } + if entry_point_result := entry_points( + group="murfey.workflows", name="data_collection_group" + ): + (workflow,) = entry_point_result + _ = workflow.load()( + message=dcg_message, + murfey_db=murfey_db, + ) + + # Register the GridSquare if it doesn't already exist + grid_square_entry = murfey_db.exec( + select(MurfeyDB.GridSquare) + .where(MurfeyDB.GridSquare.session_id == session_id) + .where(MurfeyDB.GridSquare.name == lamella_num) + .where(MurfeyDB.GridSquare.tag == dcg_tag) + ).one_or_none() + if grid_square_entry is None: + dcg_entry = murfey_db.exec( + select(MurfeyDB.DataCollectionGroup) + .where(MurfeyDB.DataCollectionGroup.session_id == session_id) + .where(MurfeyDB.DataCollectionGroup.tag == dcg_tag) + ).one() + grid_square_ispyb_result = _transport_object.do_insert_grid_square( + atlas_id=dcg_entry.atlas_id, + grid_square_id=lamella_num, + grid_square_parameters=GridSquareParameters( + tag=dcg_tag, + ), + ) + # Register to Murfey + grid_square_entry = MurfeyDB.GridSquare( + id=grid_square_ispyb_result.get("return_value", None), + name=lamella_num, + session_id=session_id, + tag=dcg_tag, + ) + murfey_db.add(grid_square_entry) + murfey_db.commit() + + return grid_square_entry + + +MILLING_STEP_LOOKUP = ( + # Match the MillingStep key names to their ISPyB IDs + # Preparation stage + ( + ( + ("eucentric_tilt", 2), + ("artificial_features", 5), + ("milling_angle", 8), + ("image_acquisition", 11), + ("lamella_placement", 14), + ), + "preparation_site", + ), + # Milling stage + ( + ( + ("delay_1", 20), + ("reference_definition", 23), + ("reference_definition_electron", 26), + ("stress_relief_cuts", 29), + ("reference_redefinition_1", 32), + ("rough_milling", 35), + ("rough_milling_electron", 38), + ("reference_redefinition_2", 41), + ("medium_milling", 44), + ("medium_milling_electron", 47), + ("fine_milling", 50), + ("fine_milling_electron", 53), + ("finer_milling", 56), + ("finer_milling_electron", 59), + ), + "chunk_site", + ), + # Thinning stage + ( + ( + ("delay_2", 65), + ("polishing_1", 68), + ("polishing_1_electron", 71), + ("polishing_2", 74), + ("polishing_2_ion", 77), + ("polishing_2_electron", 80), + ), + "thinning_site", + ), +) + + +def _register_milling_step( + site_info: LamellaSiteInfo, + grid_square: MurfeyDB.GridSquare, + murfey_db: SQLModelSession, +): + """ + Registers FIB milling metadata for the current lamella site as MillingStep entries + in ISPyB. If successful, will proceed to create a backup copy of the inserted row + in Murfey. + """ + + # Check if TransportManager has been configured + if _transport_object is None: + logger.error("No TransportManager object was configured") + return None + # Check information for milling steps is present + if site_info.steps is None: + logger.error("No milling step info found in current message") + return None + if site_info.stage_info is None: + logger.error("No stage info found in current message") + return None + + # Iteratively go through the LamellaSiteInfo model and insert for each step + for steps, stage_name in MILLING_STEP_LOOKUP: + for step_name, step_id in steps: + step_info: MillingStepInfo | None = site_info.steps.__getattribute__( + step_name + ) + if step_info is None: + continue + stage_values: StagePositionValues | None = ( + site_info.stage_info.__getattribute__(stage_name) + ) + if stage_values is None: + stage_values = StagePositionValues() + + # Check if the step has already been registered in Murfey + milling_step_entry = murfey_db.exec( + select(MurfeyDB.MillingStep) + .where(MurfeyDB.MillingStep.grid_square_id == grid_square.id) + .where(MurfeyDB.MillingStep.recipe_name == step_info.recipe_name) + .where(MurfeyDB.MillingStep.activity_name == step_info.step_name) + ).one_or_none() + + if milling_step_entry is None: + # Create a new ISPyB entry if no Murfey one is found + record = ISPyBDB.MillingStep( + millingStepNameId=step_id, + isEnabled=step_info.is_enabled, + status=step_info.status, + executionTime=step_info.execution_time, + stageX=stage_values.x, + stageY=stage_values.y, + stageZ=stage_values.z, + rotation=stage_values.rotation, + alphaTilt=stage_values.tilt_alpha, + beamType=step_info.beam_type, + beamVoltage=step_info.voltage, + beamCurrent=step_info.current, + millingAngle=step_info.milling_angle, + depthCorrection=step_info.depth_correction, + lamellaOffset=step_info.lamella_offset, + trenchHeightFront=step_info.trench_height_front, + trenchHeightRear=step_info.trench_height_rear, + widthOverlapFrontLeft=step_info.width_overlap_front_left, + widthOverlapFrontRight=step_info.width_overlap_front_right, + widthOverlapRearLeft=step_info.width_overlap_rear_left, + widthOverlapRearRight=step_info.width_overlap_rear_right, + ) + result = _transport_object.do_insert_milling_step(record) + if result.get("return_value") is None: + logger.error( + f"No MillingStep entry created for {step_info.step_name}" + ) + continue + + # Create a corresponding record in Murfey + milling_step_entry = MurfeyDB.MillingStep( + id=int(result["return_value"]), + grid_square_id=grid_square.id, + recipe_name=step_info.recipe_name, + activity_name=step_info.step_name, + is_enabled=step_info.is_enabled, + status=step_info.status, + execution_time=step_info.execution_time, + stage_x=stage_values.x, + stage_y=stage_values.y, + stage_z=stage_values.z, + rotation=stage_values.rotation, + tilt_alpha=stage_values.tilt_alpha, + beam_type=step_info.beam_type, + beam_voltage=step_info.voltage, + beam_current=step_info.current, + milling_angle=step_info.milling_angle, + depth_correction=step_info.depth_correction, + lamella_offset=step_info.lamella_offset, + trench_height_front=step_info.trench_height_front, + trench_height_rear=step_info.trench_height_rear, + width_overlap_front_left=step_info.width_overlap_front_left, + width_overlap_front_right=step_info.width_overlap_front_right, + width_overlap_rear_left=step_info.width_overlap_rear_left, + width_overlap_rear_right=step_info.width_overlap_rear_right, + ) + else: + # Update the existing ISPyB one if it already exists + _transport_object.do_update_milling_step( + milling_step_id=milling_step_entry.id, + is_enabled=step_info.is_enabled, + status=step_info.status, + execution_time=step_info.execution_time, + stage_x=stage_values.x, + stage_y=stage_values.y, + stage_z=stage_values.z, + rotation=stage_values.rotation, + tilt_alpha=stage_values.tilt_alpha, + beam_type=step_info.beam_type, + beam_voltage=step_info.voltage, + beam_current=step_info.current, + milling_angle=step_info.milling_angle, + depth_correction=step_info.depth_correction, + lamella_offset=step_info.lamella_offset, + trench_height_front=step_info.trench_height_front, + trench_height_rear=step_info.trench_height_rear, + width_overlap_front_left=step_info.width_overlap_front_left, + width_overlap_front_right=step_info.width_overlap_front_right, + width_overlap_rear_left=step_info.width_overlap_rear_left, + width_overlap_rear_right=step_info.width_overlap_rear_right, + ) + + # Update the existing Murfey one + milling_step_entry.is_enabled = step_info.is_enabled + milling_step_entry.status = step_info.status + milling_step_entry.execution_time = step_info.execution_time + milling_step_entry.stage_x = stage_values.x + milling_step_entry.stage_y = stage_values.y + milling_step_entry.stage_z = stage_values.z + milling_step_entry.rotation = stage_values.rotation + milling_step_entry.tilt_alpha = stage_values.tilt_alpha + milling_step_entry.beam_type = step_info.beam_type + milling_step_entry.beam_voltage = step_info.voltage + milling_step_entry.beam_current = step_info.current + milling_step_entry.milling_angle = step_info.milling_angle + milling_step_entry.depth_correction = step_info.depth_correction + milling_step_entry.lamella_offset = step_info.lamella_offset + milling_step_entry.trench_height_front = step_info.trench_height_front + milling_step_entry.trench_height_rear = step_info.trench_height_rear + milling_step_entry.width_overlap_front_left = ( + step_info.width_overlap_front_left + ) + milling_step_entry.width_overlap_front_right = ( + step_info.width_overlap_front_right + ) + milling_step_entry.width_overlap_rear_left = ( + step_info.width_overlap_rear_left + ) + milling_step_entry.width_overlap_rear_right = ( + step_info.width_overlap_rear_right + ) + # Mark entry for committing + murfey_db.add(milling_step_entry) + + # Commit all changes at once + murfey_db.commit() + return None + + def run(message: dict[str, Any], murfey_db: SQLModelSession): + # Outer try-finally block to handle database cleanup try: - session_id = int(message["session_id"]) - site_info = LamellaSiteInfo(**message["site_info"]) - logger.debug( - "Received the following FIB metadata for registration:\n" - f"{json.dumps(site_info.model_dump(exclude_none=True), indent=2, default=str)}" + try: + # Parse and unpack incoming message + session_id = int(message["session_id"]) + site_info = LamellaSiteInfo(**message["site_info"]) + logger.debug( + "Received the following FIB metadata for registration:\n" + f"{json.dumps(site_info.model_dump(exclude_none=True), indent=2, default=str)}" + ) + except Exception: + logger.error("Error parsing contents of message", exc_info=True) + return {"success": False, "requeue": False} + + try: + # Load instrument name and visit ID + murfey_session = murfey_db.exec( + select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id) + ).one() + visit_name = murfey_session.visit + instrument_name = murfey_session.instrument_name + except Exception: + logger.error( + "Exception encountered while querying Murfey database", exc_info=True + ) + return {"success": False, "requeue": False} + + try: + # Register the prerequisite information for this site + grid_square_entry = _ensure_prerequisites( + session_id=session_id, + instrument_name=instrument_name, + visit_name=visit_name, + site_info=site_info, + murfey_db=murfey_db, + ) + except Exception: + logger.error( + "Exception encountered while registering preqrequisite database entries", + exc_info=True, + ) + return {"success": False, "requeue": False} + if grid_square_entry is None: + logger.error( + f"Could not create GridSquare database entry for site {site_info.site_name}" + ) + return {"success": False, "requeue": False} + + try: + # Insert or update MillingStep entries + _register_milling_step( + site_info=site_info, + grid_square=grid_square_entry, + murfey_db=murfey_db, + ) + except Exception: + logger.error( + "Exception encountered while registering milling progress", + exc_info=True, + ) + return {"success": False, "requeue": False} + logger.info( + f"Successfully registered milling progress of site {site_info.site_name}" ) - except Exception: - logger.error("Error parsing contents of message", exc_info=True) - return {"success": False, "requeue": False} + return {"success": True} + finally: + murfey_db.close() From 03227a26c4cc741f311ab6b14ca7b326aa7fca71 Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Wed, 3 Jun 2026 15:22:02 +0100 Subject: [PATCH 08/19] Add field validation logic to convert stringified 'None' into Pythonic 'None' for the 'status' field --- src/murfey/util/models.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/murfey/util/models.py b/src/murfey/util/models.py index f2f00fc0e..45cf10fe1 100644 --- a/src/murfey/util/models.py +++ b/src/murfey/util/models.py @@ -179,6 +179,12 @@ class MillingStepInfo(BaseModel): width_overlap_rear_left: float | None = None width_overlap_rear_right: float | None = None + @field_validator("status", mode="before") + def handle_stringified_none(cls, v: Any) -> None: + if isinstance(v, str) and v.lower() == "none": + return None + return v + class MillingSteps(BaseModel): # Processing steps supported by AutoTEM From aba4a52b952b2254c382e0a7c293164f0f44174f Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Wed, 3 Jun 2026 15:30:20 +0100 Subject: [PATCH 09/19] Bugfixes to the MillingStep insertion logic: * GridSquare ID needs to be passed in when creating a new ISPyB record * Need to check if MillingStep update on ISPyB was successful before updating the corresponding Murfey record * Get site number from Pydantic model instead of working it out from first priciples again --- .../fib/register_milling_progress.py | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/src/murfey/workflows/fib/register_milling_progress.py b/src/murfey/workflows/fib/register_milling_progress.py index c745676bc..45a9223ca 100644 --- a/src/murfey/workflows/fib/register_milling_progress.py +++ b/src/murfey/workflows/fib/register_milling_progress.py @@ -8,7 +8,6 @@ import murfey.util.db as MurfeyDB from murfey.server import _transport_object -from murfey.util.fib import number_from_name from murfey.util.models import ( GridSquareParameters, LamellaSiteInfo, @@ -40,6 +39,10 @@ def _ensure_prerequisites( if site_info.project_name is None: logger.error("Could not construct lookup tags; 'project_name' is missing") return None + if site_info.site_number is None: + logger.error("Could not construct lookup tags; 'site_number' is missing") + return None + site_number = site_info.site_number if site_info.stage_info is None: logger.error("Could not construct lookup tags; 'stage_info' is missing") return None @@ -55,7 +58,6 @@ def _ensure_prerequisites( return None # Construct the DataCollectionGroup and GridSquare lookup tags - lamella_num = number_from_name(site_info.site_name) dcg_tag = f"{site_info.project_name}--slot_{slot_number}" # Determine variables to register data collection group and atlas with @@ -101,7 +103,7 @@ def _ensure_prerequisites( grid_square_entry = murfey_db.exec( select(MurfeyDB.GridSquare) .where(MurfeyDB.GridSquare.session_id == session_id) - .where(MurfeyDB.GridSquare.name == lamella_num) + .where(MurfeyDB.GridSquare.name == site_number) .where(MurfeyDB.GridSquare.tag == dcg_tag) ).one_or_none() if grid_square_entry is None: @@ -112,7 +114,7 @@ def _ensure_prerequisites( ).one() grid_square_ispyb_result = _transport_object.do_insert_grid_square( atlas_id=dcg_entry.atlas_id, - grid_square_id=lamella_num, + grid_square_id=site_number, grid_square_parameters=GridSquareParameters( tag=dcg_tag, ), @@ -120,7 +122,7 @@ def _ensure_prerequisites( # Register to Murfey grid_square_entry = MurfeyDB.GridSquare( id=grid_square_ispyb_result.get("return_value", None), - name=lamella_num, + name=site_number, session_id=session_id, tag=dcg_tag, ) @@ -226,7 +228,10 @@ def _register_milling_step( if milling_step_entry is None: # Create a new ISPyB entry if no Murfey one is found record = ISPyBDB.MillingStep( + # IDs millingStepNameId=step_id, + gridSquareId=grid_square.id, + # Values isEnabled=step_info.is_enabled, status=step_info.status, executionTime=step_info.execution_time, @@ -284,7 +289,7 @@ def _register_milling_step( ) else: # Update the existing ISPyB one if it already exists - _transport_object.do_update_milling_step( + result = _transport_object.do_update_milling_step( milling_step_id=milling_step_entry.id, is_enabled=step_info.is_enabled, status=step_info.status, @@ -307,6 +312,11 @@ def _register_milling_step( width_overlap_rear_left=step_info.width_overlap_rear_left, width_overlap_rear_right=step_info.width_overlap_rear_right, ) + if result.get("return_value", None) is None: + logger.error( + f"Could not update MillingStep entry for {step_info.step_name}" + ) + continue # Update the existing Murfey one milling_step_entry.is_enabled = step_info.is_enabled From 6b57da5c84a1a32cfb9c7eb8368eccebe2344c65 Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Wed, 3 Jun 2026 16:09:30 +0100 Subject: [PATCH 10/19] Forgot to add 'classmethod' decorator --- src/murfey/util/models.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/murfey/util/models.py b/src/murfey/util/models.py index 45cf10fe1..a3e594118 100644 --- a/src/murfey/util/models.py +++ b/src/murfey/util/models.py @@ -180,6 +180,7 @@ class MillingStepInfo(BaseModel): width_overlap_rear_right: float | None = None @field_validator("status", mode="before") + @classmethod def handle_stringified_none(cls, v: Any) -> None: if isinstance(v, str) and v.lower() == "none": return None From b6add3dc93cae13cec961a4e753779146586dc0f Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Wed, 3 Jun 2026 16:59:01 +0100 Subject: [PATCH 11/19] Add integrated test for the FIB milling progress workflow --- .../fib/test_register_milling_progress.py | 392 ++++++++++++++++++ 1 file changed, 392 insertions(+) create mode 100644 tests/workflows/fib/test_register_milling_progress.py diff --git a/tests/workflows/fib/test_register_milling_progress.py b/tests/workflows/fib/test_register_milling_progress.py new file mode 100644 index 000000000..13d771520 --- /dev/null +++ b/tests/workflows/fib/test_register_milling_progress.py @@ -0,0 +1,392 @@ +from pathlib import Path +from unittest.mock import MagicMock + +from pytest_mock import MockerFixture +from sqlalchemy.orm import Session as SQLAlchemySession +from sqlmodel import Session as SQLModelSession, select as sm_select + +import murfey.util.db as MurfeyDB +from murfey.workflows.fib.register_milling_progress import run +from tests.conftest import ExampleVisit + +# Module-wide variables +session_id = 10 +visit_name = f"{ExampleVisit.proposal_code}{ExampleVisit.proposal_number}-{ExampleVisit.visit_number}" +instrument_name = ExampleVisit.instrument_name + + +# Construct test FIB AutoTEM site info for reading (copied and adapted from actual session) +site_info = { + "project_name": visit_name, + "site_name": "Lamella", + "site_number": 1, + "stage_info": { + "preparation_site": { + "x": -0.0031091224743875403, + "y": 0.00420867925495798, + "z": 0.0323644854106331, + "rotation": 285.003247202109, + "tilt_alpha": 25.9996646026832, + "slot_number": 1, + }, + "chunk_site": { + "x": -0.0030037500000000003, + "y": 0.004293, + "z": 0.032350405092592606, + "rotation": 285.003247202109, + "tilt_alpha": -0.000134158926728586, + "slot_number": 1, + }, + "thinning_site": { + "x": -0.0030037500000000003, + "y": 0.004293, + "z": 0.032350405092592606, + "rotation": 285.003247202109, + "tilt_alpha": -0.000134158926728586, + "slot_number": 1, + }, + "chunk_coincidence_params": { + "x": -0.0030048260286678298, + "y": 0.004308828126160981, + "z": 0.0323400707790533, + "rotation": 285.003247202109, + "tilt_alpha": -0.000134158926728586, + "slot_number": 1, + }, + "thinning_params": { + "x": -0.0030037500000000003, + "y": 0.004293, + "z": 0.032350405092592606, + "rotation": 285.003247202109, + "tilt_alpha": -0.000134158926728586, + "slot_number": 1, + }, + }, + "steps": { + "eucentric_tilt": { + "step_name": "Eucentric Tilt", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "Finished", + "execution_time": 328.6657458, + }, + "artificial_features": { + "step_name": "Artificial Features", + "recipe_name": "Preparation", + "is_enabled": False, + "execution_time": 0.0, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 5e-10, + }, + "milling_angle": { + "step_name": "Milling Angle", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "Finished", + "execution_time": 208.4942922, + "milling_angle": 12.0, + }, + "image_acquisition": { + "step_name": "Image Acquisition", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "Finished", + "execution_time": 19.7872887, + "site_location_type": "Chunk", + }, + "lamella_placement": { + "step_name": "Lamella Placement", + "recipe_name": "Preparation", + "is_enabled": False, + "execution_time": 0.0, + }, + "delay_1": { + "step_name": "Delay", + "recipe_name": "Milling", + "is_enabled": False, + "execution_time": 0.0, + }, + "reference_definition": { + "step_name": "Reference Definition", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 75.6320442, + "site_location_type": "Chunk", + }, + "reference_definition_electron": { + "step_name": "Electron Reference Definition", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 131.7462301, + }, + "stress_relief_cuts": { + "step_name": "Stress Relief Cuts", + "recipe_name": "Milling", + "is_enabled": False, + "execution_time": 0.0, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 1e-09, + "depth_correction": 3.0, + }, + "reference_redefinition_1": { + "step_name": "Reference Redefinition 1", + "recipe_name": "Milling", + "is_enabled": False, + "execution_time": 0.0, + "site_location_type": "Chunk", + }, + "rough_milling": { + "step_name": "Rough Milling", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 1.5929074719, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 1e-09, + "depth_correction": 3.0, + "lamella_offset": 2e-06, + "trench_height_front": 5.6338138462765e-06, + "trench_height_rear": 8.60473362403628e-06, + "width_overlap_front_left": 2e-06, + "width_overlap_front_right": 2e-06, + "width_overlap_rear_left": 2e-06, + "width_overlap_rear_right": 2e-06, + }, + "rough_milling_electron": { + "step_name": "Rough Milling - Electron Image", + "recipe_name": "Milling", + "is_enabled": False, + "execution_time": 0.0, + "site_location_type": "Chunk", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 1.25e-11, + }, + "reference_redefinition_2": { + "step_name": "Reference Redefinition 2", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 105.1459769, + "site_location_type": "Chunk", + }, + "medium_milling": { + "step_name": "Medium Milling", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 298.1471377, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 5e-10, + "depth_correction": 2.0, + "lamella_offset": 1.5e-06, + "width_overlap_front_left": 1.5e-06, + "width_overlap_front_right": 1.5e-06, + "width_overlap_rear_left": 1.5e-06, + "width_overlap_rear_right": 1.5e-06, + }, + "medium_milling_electron": { + "step_name": "Medium Milling - Electron Image", + "recipe_name": "Milling", + "is_enabled": False, + "execution_time": 0.0, + "site_location_type": "Chunk", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 1.25e-11, + }, + "fine_milling": { + "step_name": "Fine Milling", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 397.414388, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 3e-10, + "depth_correction": 2.0, + "lamella_offset": 1e-06, + "width_overlap_front_left": 1e-06, + "width_overlap_front_right": 1e-06, + "width_overlap_rear_left": 1e-06, + "width_overlap_rear_right": 1e-06, + }, + "fine_milling_electron": { + "step_name": "Fine Milling - Electron Image", + "recipe_name": "Milling", + "is_enabled": False, + "execution_time": 0.0, + "site_location_type": "Chunk", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 1.25e-11, + }, + "finer_milling": { + "step_name": "Finer Milling", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 887.9686893, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 1e-10, + "depth_correction": 2.0, + "lamella_offset": 4.0000000000000003e-07, + "width_overlap_front_left": 5.000000000000001e-07, + "width_overlap_front_right": 5.000000000000001e-07, + "width_overlap_rear_left": 5.000000000000001e-07, + "width_overlap_rear_right": 5.000000000000001e-07, + }, + "finer_milling_electron": { + "step_name": "Finer Milling - Electron Image", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 59.2296577, + "site_location_type": "Chunk", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 2.5e-11, + }, + "delay_2": { + "step_name": "Delay", + "recipe_name": "Thinning", + "is_enabled": False, + "execution_time": 0.0, + }, + "polishing_1": { + "step_name": "Polishing 1", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "Finished", + "execution_time": 680.5567315, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 5e-11, + "depth_correction": 2.0, + "lamella_offset": 2.5000000000000004e-07, + "width_overlap_front_left": 0.0, + "width_overlap_front_right": 0.0, + "width_overlap_rear_left": 0.0, + "width_overlap_rear_right": 0.0, + }, + "polishing_1_electron": { + "step_name": "Polishing 1 - Electron Image", + "recipe_name": "Thinning", + "is_enabled": False, + "execution_time": 0.0, + "site_location_type": "Thinning", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 1.25e-11, + }, + "polishing_2": { + "step_name": "Polishing 2", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "Finished", + "execution_time": 1.170488927, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 3e-11, + "depth_correction": 1.5, + "lamella_offset": 0.0, + "width_overlap_front_left": 0.0, + "width_overlap_front_right": 0.0, + "width_overlap_rear_left": 0.0, + "width_overlap_rear_right": 0.0, + }, + "polishing_2_ion": { + "step_name": "Polishing 2 - Ion Image", + "recipe_name": "Thinning", + "is_enabled": False, + "execution_time": 0.0, + "site_location_type": "Thinning", + "beam_type": "Ion", + "voltage": 30000.0, + "current": 3e-11, + }, + "polishing_2_electron": { + "step_name": "Polishing 2 - Electron Image", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "Finished", + "execution_time": 56.9180832, + "site_location_type": "Thinning", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 2.5e-11, + }, + }, +} +# Construct the RabbitMQ message received +message = { + "register": "fib.register_milling_progress", + "session_id": session_id, + "site_info": site_info, +} + + +def test_run_with_db( + mocker: MockerFixture, + murfey_db_session: SQLModelSession, + ispyb_db_session: SQLAlchemySession, + mock_ispyb_credentials: Path, +): + # Add a test visit to the database + if not ( + session_entry := murfey_db_session.exec( + sm_select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id) + ).one_or_none() + ): + session_entry = MurfeyDB.Session(id=session_id) + session_entry.name = visit_name + session_entry.visit = visit_name + session_entry.instrument_name = instrument_name + + murfey_db_session.add(session_entry) + murfey_db_session.commit() + + # Mock the ISPyB connection where the TransportManager class is located + mock_security_config = MagicMock() + mock_security_config.ispyb_credentials = mock_ispyb_credentials + mocker.patch( + "murfey.server.ispyb.get_security_config", + return_value=mock_security_config, + ) + mocker.patch( + "murfey.server.ispyb.ISPyBSession", + return_value=ispyb_db_session, + ) + + # Mock the ISPYB connection when registering DataCollectionGroup + mocker.patch( + "murfey.workflows.register_data_collection_group.ISPyBSession", + return_value=ispyb_db_session, + ) + + # Patch the TransportManager object in the workflows called + from murfey.server.ispyb import TransportManager + + mocker.patch( + "murfey.workflows.register_data_collection_group._transport_object", + new=TransportManager("PikaTransport"), + ) + mocker.patch( + "murfey.workflows.fib.register_milling_progress._transport_object", + new=TransportManager("PikaTransport"), + ) + + # Run the workflow twice (fresh insert and update existing) + for i in range(2): + result = run( + message=message, + murfey_db=murfey_db_session, + ) + assert result.get("success", False) From f57a1f5be1644b7deea611773c5fb3564bb07e2c Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Wed, 3 Jun 2026 17:10:11 +0100 Subject: [PATCH 12/19] Updated ISPyB schema version --- .github/workflows/test.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4f34779c4..bcf26da33 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -3,10 +3,11 @@ name: Build and test on: [push, pull_request] env: - ISPYB_DATABASE_SCHEMA: 4.11.0 + ISPYB_DATABASE_SCHEMA: 5.1.0 # Installs from GitHub # Versions: https://github.com/DiamondLightSource/ispyb-database/tags # Previous version(s): + # 4.11.0 # 4.8.0 # 4.2.1 # released 2024-08-19 # 4.1.0 # released 2024-03-26 From 1218721986e0e148f24b696a70c6370b701ccfb4 Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Wed, 3 Jun 2026 17:56:03 +0100 Subject: [PATCH 13/19] Move 'MillingStep' record creation into 'do_insert_milling_step' instead, since the 'MillingStepNameId' requires a DB query --- src/murfey/server/ispyb.py | 65 ++++++++++++++++++- .../fib/register_milling_progress.py | 60 ++++++++++------- 2 files changed, 100 insertions(+), 25 deletions(-) diff --git a/src/murfey/server/ispyb.py b/src/murfey/server/ispyb.py index 310688106..75525dd0c 100644 --- a/src/murfey/server/ispyb.py +++ b/src/murfey/server/ispyb.py @@ -21,6 +21,7 @@ FoilHole, GridSquare, MillingStep, + MillingStepName, ProcessingJob, ProcessingJobParameter, Proposal, @@ -711,9 +712,71 @@ def do_update_processing_status(self, record: AutoProcProgram, **kwargs): ) return {"success": False, "return_value": None} - def do_insert_milling_step(self, record: MillingStep): + def do_insert_milling_step( + self, + # MillingStepName identifiers + recipe_name: str, + activity_name: str, + grid_square_id: int, + # Values + is_enabled: bool | None = None, + status: str | None = None, + execution_time: float | None = None, + stage_x: float | None = None, + stage_y: float | None = None, + stage_z: float | None = None, + rotation: float | None = None, + tilt_alpha: float | None = None, + beam_type: str | None = None, + beam_voltage: float | None = None, + beam_current: float | None = None, + milling_angle: float | None = None, + depth_correction: float | None = None, + lamella_offset: float | None = None, + trench_height_front: float | None = None, + trench_height_rear: float | None = None, + width_overlap_front_left: float | None = None, + width_overlap_front_right: float | None = None, + width_overlap_rear_left: float | None = None, + width_overlap_rear_right: float | None = None, + ): try: with ISPyBSession() as db: + # Find the ID of this MillingStep + milling_step_name = ( + db.query(MillingStepName) + .filter( + MillingStepName.recipe == recipe_name, + MillingStepName.step == activity_name, + ) + .one() + ) + record = MillingStep( + # IDs + millingStepNameId=milling_step_name.millingStepNameId, + gridSquareId=grid_square_id, + # Values + isEnabled=is_enabled, + status=status, + executionTime=execution_time, + stageX=stage_x, + stageY=stage_y, + stageZ=stage_z, + rotation=rotation, + alphaTilt=tilt_alpha, + beamType=beam_type, + beamVoltage=beam_voltage, + beamCurrent=beam_current, + millingAngle=milling_angle, + depthCorrection=depth_correction, + lamellaOffset=lamella_offset, + trenchHeightFront=trench_height_front, + trenchHeightRear=trench_height_rear, + widthOverlapFrontLeft=width_overlap_front_left, + widthOverlapFrontRight=width_overlap_front_right, + widthOverlapRearLeft=width_overlap_rear_left, + widthOverlapRearRight=width_overlap_rear_right, + ) db.add(record) db.commit() log.info(f"Created MillingStep {record.millingStepId}") diff --git a/src/murfey/workflows/fib/register_milling_progress.py b/src/murfey/workflows/fib/register_milling_progress.py index 45a9223ca..dc3562eb5 100644 --- a/src/murfey/workflows/fib/register_milling_progress.py +++ b/src/murfey/workflows/fib/register_milling_progress.py @@ -3,7 +3,6 @@ from importlib.metadata import entry_points from typing import Any -import ispyb.sqlalchemy._auto_db_schema as ISPyBDB from sqlmodel import Session as SQLModelSession, select import murfey.util.db as MurfeyDB @@ -202,15 +201,28 @@ def _register_milling_step( if site_info.stage_info is None: logger.error("No stage info found in current message") return None + # Check that GridSquare has ID (for type checking) + if grid_square.id is None: + logger.error("Current GridSquare entry has no ID") + return None # Iteratively go through the LamellaSiteInfo model and insert for each step for steps, stage_name in MILLING_STEP_LOOKUP: - for step_name, step_id in steps: + for step_name, _ in steps: step_info: MillingStepInfo | None = site_info.steps.__getattribute__( step_name ) + # Early continues if key information is missing if step_info is None: + logger.debug(f"No step info found for {step_name}") + continue + if step_info.recipe_name is None: + logger.debug(f"No recipe name found for {step_name}") continue + if step_info.step_name is None: + logger.debug(f"No step name found for {step_name}") + continue + stage_values: StagePositionValues | None = ( site_info.stage_info.__getattribute__(stage_name) ) @@ -227,33 +239,33 @@ def _register_milling_step( if milling_step_entry is None: # Create a new ISPyB entry if no Murfey one is found - record = ISPyBDB.MillingStep( + result = _transport_object.do_insert_milling_step( # IDs - millingStepNameId=step_id, - gridSquareId=grid_square.id, + recipe_name=step_info.recipe_name, + activity_name=step_info.step_name, + grid_square_id=grid_square.id, # Values - isEnabled=step_info.is_enabled, + is_enabled=step_info.is_enabled, status=step_info.status, - executionTime=step_info.execution_time, - stageX=stage_values.x, - stageY=stage_values.y, - stageZ=stage_values.z, + execution_time=step_info.execution_time, + stage_x=stage_values.x, + stage_y=stage_values.y, + stage_z=stage_values.z, rotation=stage_values.rotation, - alphaTilt=stage_values.tilt_alpha, - beamType=step_info.beam_type, - beamVoltage=step_info.voltage, - beamCurrent=step_info.current, - millingAngle=step_info.milling_angle, - depthCorrection=step_info.depth_correction, - lamellaOffset=step_info.lamella_offset, - trenchHeightFront=step_info.trench_height_front, - trenchHeightRear=step_info.trench_height_rear, - widthOverlapFrontLeft=step_info.width_overlap_front_left, - widthOverlapFrontRight=step_info.width_overlap_front_right, - widthOverlapRearLeft=step_info.width_overlap_rear_left, - widthOverlapRearRight=step_info.width_overlap_rear_right, + tilt_alpha=stage_values.tilt_alpha, + beam_type=step_info.beam_type, + beam_voltage=step_info.voltage, + beam_current=step_info.current, + milling_angle=step_info.milling_angle, + depth_correction=step_info.depth_correction, + lamella_offset=step_info.lamella_offset, + trench_height_front=step_info.trench_height_front, + trench_height_rear=step_info.trench_height_rear, + width_overlap_front_left=step_info.width_overlap_front_left, + width_overlap_front_right=step_info.width_overlap_front_right, + width_overlap_rear_left=step_info.width_overlap_rear_left, + width_overlap_rear_right=step_info.width_overlap_rear_right, ) - result = _transport_object.do_insert_milling_step(record) if result.get("return_value") is None: logger.error( f"No MillingStep entry created for {step_info.step_name}" From 82169a5d6e231ac2b3171c4883fba6abae90ada8 Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Wed, 3 Jun 2026 18:00:34 +0100 Subject: [PATCH 14/19] Simplify lookup logic --- .../fib/register_milling_progress.py | 52 +++++++++---------- 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/src/murfey/workflows/fib/register_milling_progress.py b/src/murfey/workflows/fib/register_milling_progress.py index dc3562eb5..c12d47840 100644 --- a/src/murfey/workflows/fib/register_milling_progress.py +++ b/src/murfey/workflows/fib/register_milling_progress.py @@ -136,43 +136,43 @@ def _ensure_prerequisites( # Preparation stage ( ( - ("eucentric_tilt", 2), - ("artificial_features", 5), - ("milling_angle", 8), - ("image_acquisition", 11), - ("lamella_placement", 14), + "eucentric_tilt", + "artificial_features", + "milling_angle", + "image_acquisition", + "lamella_placement", ), "preparation_site", ), # Milling stage ( ( - ("delay_1", 20), - ("reference_definition", 23), - ("reference_definition_electron", 26), - ("stress_relief_cuts", 29), - ("reference_redefinition_1", 32), - ("rough_milling", 35), - ("rough_milling_electron", 38), - ("reference_redefinition_2", 41), - ("medium_milling", 44), - ("medium_milling_electron", 47), - ("fine_milling", 50), - ("fine_milling_electron", 53), - ("finer_milling", 56), - ("finer_milling_electron", 59), + "delay_1", + "reference_definition", + "reference_definition_electron", + "stress_relief_cuts", + "reference_redefinition_1", + "rough_milling", + "rough_milling_electron", + "reference_redefinition_2", + "medium_milling", + "medium_milling_electron", + "fine_milling", + "fine_milling_electron", + "finer_milling", + "finer_milling_electron", ), "chunk_site", ), # Thinning stage ( ( - ("delay_2", 65), - ("polishing_1", 68), - ("polishing_1_electron", 71), - ("polishing_2", 74), - ("polishing_2_ion", 77), - ("polishing_2_electron", 80), + "delay_2", + "polishing_1", + "polishing_1_electron", + "polishing_2", + "polishing_2_ion", + "polishing_2_electron", ), "thinning_site", ), @@ -208,7 +208,7 @@ def _register_milling_step( # Iteratively go through the LamellaSiteInfo model and insert for each step for steps, stage_name in MILLING_STEP_LOOKUP: - for step_name, _ in steps: + for step_name in steps: step_info: MillingStepInfo | None = site_info.steps.__getattribute__( step_name ) From 6818e08995d7e3144ebafa66cef2646e0e671bae Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Wed, 3 Jun 2026 18:29:15 +0100 Subject: [PATCH 15/19] Add stringified 'None' values as statuses --- tests/workflows/fib/test_register_milling_progress.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/workflows/fib/test_register_milling_progress.py b/tests/workflows/fib/test_register_milling_progress.py index 13d771520..5eedcb448 100644 --- a/tests/workflows/fib/test_register_milling_progress.py +++ b/tests/workflows/fib/test_register_milling_progress.py @@ -74,6 +74,7 @@ "step_name": "Artificial Features", "recipe_name": "Preparation", "is_enabled": False, + "status": "None", "execution_time": 0.0, "beam_type": "Ion", "voltage": 30000.0, @@ -99,12 +100,14 @@ "step_name": "Lamella Placement", "recipe_name": "Preparation", "is_enabled": False, + "status": "None", "execution_time": 0.0, }, "delay_1": { "step_name": "Delay", "recipe_name": "Milling", "is_enabled": False, + "status": "None", "execution_time": 0.0, }, "reference_definition": { @@ -126,6 +129,7 @@ "step_name": "Stress Relief Cuts", "recipe_name": "Milling", "is_enabled": False, + "status": "None", "execution_time": 0.0, "beam_type": "Ion", "voltage": 30000.0, @@ -136,6 +140,7 @@ "step_name": "Reference Redefinition 1", "recipe_name": "Milling", "is_enabled": False, + "status": "None", "execution_time": 0.0, "site_location_type": "Chunk", }, @@ -161,6 +166,7 @@ "step_name": "Rough Milling - Electron Image", "recipe_name": "Milling", "is_enabled": False, + "status": "None", "execution_time": 0.0, "site_location_type": "Chunk", "beam_type": "Electron", @@ -195,6 +201,7 @@ "step_name": "Medium Milling - Electron Image", "recipe_name": "Milling", "is_enabled": False, + "status": "None", "execution_time": 0.0, "site_location_type": "Chunk", "beam_type": "Electron", @@ -221,6 +228,7 @@ "step_name": "Fine Milling - Electron Image", "recipe_name": "Milling", "is_enabled": False, + "status": "None", "execution_time": 0.0, "site_location_type": "Chunk", "beam_type": "Electron", @@ -258,6 +266,7 @@ "step_name": "Delay", "recipe_name": "Thinning", "is_enabled": False, + "status": "None", "execution_time": 0.0, }, "polishing_1": { @@ -280,6 +289,7 @@ "step_name": "Polishing 1 - Electron Image", "recipe_name": "Thinning", "is_enabled": False, + "status": "None", "execution_time": 0.0, "site_location_type": "Thinning", "beam_type": "Electron", @@ -306,6 +316,7 @@ "step_name": "Polishing 2 - Ion Image", "recipe_name": "Thinning", "is_enabled": False, + "status": "None", "execution_time": 0.0, "site_location_type": "Thinning", "beam_type": "Ion", From 6047bd24003706664bedd8626ece8d8b841c1f54 Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Thu, 4 Jun 2026 09:42:52 +0100 Subject: [PATCH 16/19] Centralise early exits under main 'run' function to make testing simpler --- .../fib/register_milling_progress.py | 151 ++++++++++-------- 1 file changed, 84 insertions(+), 67 deletions(-) diff --git a/src/murfey/workflows/fib/register_milling_progress.py b/src/murfey/workflows/fib/register_milling_progress.py index c12d47840..aa67b1fc7 100644 --- a/src/murfey/workflows/fib/register_milling_progress.py +++ b/src/murfey/workflows/fib/register_milling_progress.py @@ -1,7 +1,9 @@ +from __future__ import annotations + import json import logging from importlib.metadata import entry_points -from typing import Any +from typing import TYPE_CHECKING, Any from sqlmodel import Session as SQLModelSession, select @@ -11,9 +13,15 @@ GridSquareParameters, LamellaSiteInfo, MillingStepInfo, + MillingSteps, + StagePositionInfo, StagePositionValues, ) +if TYPE_CHECKING: + from murfey.server.ispyb import TransportManager + + logger = logging.getLogger("murfey.workflows.fib.register_milling_progress") @@ -21,7 +29,10 @@ def _ensure_prerequisites( session_id: int, instrument_name: str, visit_name: str, - site_info: LamellaSiteInfo, + project_name: str, + slot_number: int, + site_number: int, + transport_object: TransportManager, murfey_db: SQLModelSession, ): """ @@ -29,35 +40,8 @@ def _ensure_prerequisites( Atlas, and GridSquare placeholders in Murfey if they don't already exist """ - # Skip this step if no TransportManager object is configured - if _transport_object is None: - logger.error("No TransportManager object was configured") - return None - - # Early exits if information needed to construct lookup tags are missing - if site_info.project_name is None: - logger.error("Could not construct lookup tags; 'project_name' is missing") - return None - if site_info.site_number is None: - logger.error("Could not construct lookup tags; 'site_number' is missing") - return None - site_number = site_info.site_number - if site_info.stage_info is None: - logger.error("Could not construct lookup tags; 'stage_info' is missing") - return None - if site_info.stage_info.preparation_site is None: - logger.error("Could not construct lookup tags; 'preparation_site' is missing") - return None - if site_info.stage_info.preparation_site.slot_number is None: - logger.error("Could not construct lookup tags; 'slot_number' is missing") - return None - slot_number = site_info.stage_info.preparation_site.slot_number - if site_info.site_name is None: - logger.error("Could not construct lookup tags; 'site_name' is missing") - return None - # Construct the DataCollectionGroup and GridSquare lookup tags - dcg_tag = f"{site_info.project_name}--slot_{slot_number}" + dcg_tag = f"{project_name}--slot_{slot_number}" # Determine variables to register data collection group and atlas with proposal_code = "".join(char for char in visit_name.split("-")[0] if char.isalpha()) @@ -111,7 +95,7 @@ def _ensure_prerequisites( .where(MurfeyDB.DataCollectionGroup.session_id == session_id) .where(MurfeyDB.DataCollectionGroup.tag == dcg_tag) ).one() - grid_square_ispyb_result = _transport_object.do_insert_grid_square( + grid_square_ispyb_result = transport_object.do_insert_grid_square( atlas_id=dcg_entry.atlas_id, grid_square_id=site_number, grid_square_parameters=GridSquareParameters( @@ -180,8 +164,10 @@ def _ensure_prerequisites( def _register_milling_step( - site_info: LamellaSiteInfo, + milling_steps: MillingSteps, + stage_info: StagePositionInfo, grid_square: MurfeyDB.GridSquare, + transport_object: TransportManager, murfey_db: SQLModelSession, ): """ @@ -189,18 +175,6 @@ def _register_milling_step( in ISPyB. If successful, will proceed to create a backup copy of the inserted row in Murfey. """ - - # Check if TransportManager has been configured - if _transport_object is None: - logger.error("No TransportManager object was configured") - return None - # Check information for milling steps is present - if site_info.steps is None: - logger.error("No milling step info found in current message") - return None - if site_info.stage_info is None: - logger.error("No stage info found in current message") - return None # Check that GridSquare has ID (for type checking) if grid_square.id is None: logger.error("Current GridSquare entry has no ID") @@ -209,7 +183,7 @@ def _register_milling_step( # Iteratively go through the LamellaSiteInfo model and insert for each step for steps, stage_name in MILLING_STEP_LOOKUP: for step_name in steps: - step_info: MillingStepInfo | None = site_info.steps.__getattribute__( + step_info: MillingStepInfo | None = milling_steps.__getattribute__( step_name ) # Early continues if key information is missing @@ -223,8 +197,8 @@ def _register_milling_step( logger.debug(f"No step name found for {step_name}") continue - stage_values: StagePositionValues | None = ( - site_info.stage_info.__getattribute__(stage_name) + stage_values: StagePositionValues | None = stage_info.__getattribute__( + stage_name ) if stage_values is None: stage_values = StagePositionValues() @@ -239,7 +213,7 @@ def _register_milling_step( if milling_step_entry is None: # Create a new ISPyB entry if no Murfey one is found - result = _transport_object.do_insert_milling_step( + result = transport_object.do_insert_milling_step( # IDs recipe_name=step_info.recipe_name, activity_name=step_info.step_name, @@ -301,7 +275,7 @@ def _register_milling_step( ) else: # Update the existing ISPyB one if it already exists - result = _transport_object.do_update_milling_step( + result = transport_object.do_update_milling_step( milling_step_id=milling_step_entry.id, is_enabled=step_info.is_enabled, status=step_info.status, @@ -368,20 +342,60 @@ def _register_milling_step( def run(message: dict[str, Any], murfey_db: SQLModelSession): - # Outer try-finally block to handle database cleanup + # Early exit if no TransportManager was set up + if _transport_object is None: + logger.error("No TransportManager object was configured") + return {"success": False, "requeue": False} + try: - try: - # Parse and unpack incoming message - session_id = int(message["session_id"]) - site_info = LamellaSiteInfo(**message["site_info"]) - logger.debug( - "Received the following FIB metadata for registration:\n" - f"{json.dumps(site_info.model_dump(exclude_none=True), indent=2, default=str)}" - ) - except Exception: - logger.error("Error parsing contents of message", exc_info=True) - return {"success": False, "requeue": False} + # Parse and unpack incoming message + session_id = int(message["session_id"]) + site_info = LamellaSiteInfo(**message["site_info"]) + logger.debug( + "Received the following FIB metadata for registration:\n" + f"{json.dumps(site_info.model_dump(exclude_none=True), indent=2, default=str)}" + ) + except Exception: + logger.error("Error parsing contents of message", exc_info=True) + return {"success": False, "requeue": False} + + # Early exits if information needed to construct lookup tags are missing + # Project and site values + if site_info.project_name is None: + logger.error("Could not construct lookup tags; 'project_name' is missing") + return {"success": False, "requeue": False} + project_name = site_info.project_name + if site_info.site_number is None: + logger.error("Could not construct lookup tags; 'site_number' is missing") + return {"success": False, "requeue": False} + site_number = site_info.site_number + if site_info.site_name is None: + logger.error("Could not construct lookup tags; 'site_name' is missing") + return {"success": False, "requeue": False} + site_name = site_info.site_name + + # Stage information + if site_info.stage_info is None: + logger.error("Could not construct lookup tags; 'stage_info' is missing") + return {"success": False, "requeue": False} + stage_info = site_info.stage_info + if stage_info.preparation_site is None: + logger.error("Could not construct lookup tags; 'preparation_site' is missing") + return {"success": False, "requeue": False} + preparation_site = stage_info.preparation_site + if preparation_site.slot_number is None: + logger.error("Could not construct lookup tags; 'slot_number' is missing") + return {"success": False, "requeue": False} + slot_number = preparation_site.slot_number + + # Milling step information + if site_info.steps is None: + logger.error("No milling step info found in current message") + return None + milling_steps = site_info.steps + # Outer try-finally block to handle database cleanup + try: try: # Load instrument name and visit ID murfey_session = murfey_db.exec( @@ -401,7 +415,10 @@ def run(message: dict[str, Any], murfey_db: SQLModelSession): session_id=session_id, instrument_name=instrument_name, visit_name=visit_name, - site_info=site_info, + project_name=project_name, + slot_number=slot_number, + site_number=site_number, + transport_object=_transport_object, murfey_db=murfey_db, ) except Exception: @@ -412,15 +429,17 @@ def run(message: dict[str, Any], murfey_db: SQLModelSession): return {"success": False, "requeue": False} if grid_square_entry is None: logger.error( - f"Could not create GridSquare database entry for site {site_info.site_name}" + f"Could not create GridSquare database entry for site {site_name}" ) return {"success": False, "requeue": False} try: # Insert or update MillingStep entries _register_milling_step( - site_info=site_info, + milling_steps=milling_steps, + stage_info=stage_info, grid_square=grid_square_entry, + transport_object=_transport_object, murfey_db=murfey_db, ) except Exception: @@ -429,9 +448,7 @@ def run(message: dict[str, Any], murfey_db: SQLModelSession): exc_info=True, ) return {"success": False, "requeue": False} - logger.info( - f"Successfully registered milling progress of site {site_info.site_name}" - ) + logger.info(f"Successfully registered milling progress of site {site_name}") return {"success": True} finally: murfey_db.close() From eb5b4f9a89e6a3e1bd2a63d6fb85f4a99e1dde09 Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Thu, 4 Jun 2026 10:20:49 +0100 Subject: [PATCH 17/19] More stringent checks on the registered database results --- .../fib/test_register_milling_progress.py | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/tests/workflows/fib/test_register_milling_progress.py b/tests/workflows/fib/test_register_milling_progress.py index 5eedcb448..dd8741ddb 100644 --- a/tests/workflows/fib/test_register_milling_progress.py +++ b/tests/workflows/fib/test_register_milling_progress.py @@ -1,7 +1,10 @@ from pathlib import Path +from typing import Any, cast from unittest.mock import MagicMock +import ispyb.sqlalchemy._auto_db_schema as ISPyBDB from pytest_mock import MockerFixture +from sqlalchemy import select as sa_select from sqlalchemy.orm import Session as SQLAlchemySession from sqlmodel import Session as SQLModelSession, select as sm_select @@ -401,3 +404,64 @@ def test_run_with_db( murfey_db=murfey_db_session, ) assert result.get("success", False) + + # There should be a DataCollectionGroup entry in Murfey + dcg_murfey = murfey_db_session.exec( + sm_select(MurfeyDB.DataCollectionGroup) + .where(MurfeyDB.DataCollectionGroup.session_id == session_id) + .where( + MurfeyDB.DataCollectionGroup.tag == f"{site_info['project_name']}--slot_1" + ) + ).one_or_none() + assert dcg_murfey is not None + + # There should be a DataCollectionGroup entry in ISPyB + dcg_ispyb = ispyb_db_session.execute( + sa_select(ISPyBDB.DataCollectionGroup).where( + ISPyBDB.DataCollectionGroup.dataCollectionGroupId == dcg_murfey.id + ) + ).scalar_one_or_none() + assert dcg_ispyb is not None + + # There should be an Atlas in ISPyB + atlas_ispyb = ispyb_db_session.execute( + sa_select(ISPyBDB.Atlas).where( + ISPyBDB.Atlas.dataCollectionGroupId == dcg_ispyb.dataCollectionGroupId + ) + ).scalar_one_or_none() + assert atlas_ispyb is not None + + # There should be one GridSquare entry in ISPyB per lamella site tested + gs_ispyb_rows = ( + ispyb_db_session.execute( + sa_select(ISPyBDB.GridSquare).where( + ISPyBDB.GridSquare.atlasId == atlas_ispyb.atlasId + ) + ) + .scalars() + .all() + ) + assert len(gs_ispyb_rows) >= 1 + gs_ispyb = gs_ispyb_rows[0] + + steps = cast(dict[str, Any], site_info["steps"]) + + # There should be one MillingStep entry in ISPyB for each step in the message + milling_step_ispyb_rows = ( + ispyb_db_session.execute( + sa_select(ISPyBDB.MillingStep).where( + ISPyBDB.MillingStep.gridSquareId == gs_ispyb.gridSquareId + ) + ) + .scalars() + .all() + ) + assert len(milling_step_ispyb_rows) == len(steps.keys()) + + # There should be the same thing in Murfey + milling_step_murfey_rows = murfey_db_session.exec( + sm_select(MurfeyDB.MillingStep).where( + MurfeyDB.MillingStep.grid_square_id == gs_ispyb.gridSquareId + ) + ).all() + assert len(milling_step_murfey_rows) == len(steps.keys()) From bde424646309f13d4284912f52f2e733c779a7fe Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Thu, 4 Jun 2026 12:49:54 +0100 Subject: [PATCH 18/19] Missed converting a return statement to a dictionary --- src/murfey/workflows/fib/register_milling_progress.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/murfey/workflows/fib/register_milling_progress.py b/src/murfey/workflows/fib/register_milling_progress.py index aa67b1fc7..3c209dc38 100644 --- a/src/murfey/workflows/fib/register_milling_progress.py +++ b/src/murfey/workflows/fib/register_milling_progress.py @@ -391,7 +391,7 @@ def run(message: dict[str, Any], murfey_db: SQLModelSession): # Milling step information if site_info.steps is None: logger.error("No milling step info found in current message") - return None + return {"success": False, "requeue": False} milling_steps = site_info.steps # Outer try-finally block to handle database cleanup From b9ed626c0beb5e48e861f5ea345d9d641c240195 Mon Sep 17 00:00:00 2001 From: Eu Pin Tien Date: Thu, 4 Jun 2026 13:11:46 +0100 Subject: [PATCH 19/19] Unpacked 'milling_steps' and 'stage_info' sections from the 'site_info' master dictionary; added test function to check the early exit cases --- .../fib/test_register_milling_progress.py | 712 ++++++++++-------- 1 file changed, 392 insertions(+), 320 deletions(-) diff --git a/tests/workflows/fib/test_register_milling_progress.py b/tests/workflows/fib/test_register_milling_progress.py index dd8741ddb..606f5ba79 100644 --- a/tests/workflows/fib/test_register_milling_progress.py +++ b/tests/workflows/fib/test_register_milling_progress.py @@ -1,8 +1,8 @@ from pathlib import Path -from typing import Any, cast from unittest.mock import MagicMock import ispyb.sqlalchemy._auto_db_schema as ISPyBDB +import pytest from pytest_mock import MockerFixture from sqlalchemy import select as sa_select from sqlalchemy.orm import Session as SQLAlchemySession @@ -19,325 +19,327 @@ # Construct test FIB AutoTEM site info for reading (copied and adapted from actual session) +milling_steps = { + "eucentric_tilt": { + "step_name": "Eucentric Tilt", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "Finished", + "execution_time": 328.6657458, + }, + "artificial_features": { + "step_name": "Artificial Features", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 5e-10, + }, + "milling_angle": { + "step_name": "Milling Angle", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "Finished", + "execution_time": 208.4942922, + "milling_angle": 12.0, + }, + "image_acquisition": { + "step_name": "Image Acquisition", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "Finished", + "execution_time": 19.7872887, + "site_location_type": "Chunk", + }, + "lamella_placement": { + "step_name": "Lamella Placement", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + }, + "delay_1": { + "step_name": "Delay", + "recipe_name": "Milling", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + }, + "reference_definition": { + "step_name": "Reference Definition", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 75.6320442, + "site_location_type": "Chunk", + }, + "reference_definition_electron": { + "step_name": "Electron Reference Definition", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 131.7462301, + }, + "stress_relief_cuts": { + "step_name": "Stress Relief Cuts", + "recipe_name": "Milling", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 1e-09, + "depth_correction": 3.0, + }, + "reference_redefinition_1": { + "step_name": "Reference Redefinition 1", + "recipe_name": "Milling", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "site_location_type": "Chunk", + }, + "rough_milling": { + "step_name": "Rough Milling", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 1.5929074719, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 1e-09, + "depth_correction": 3.0, + "lamella_offset": 2e-06, + "trench_height_front": 5.6338138462765e-06, + "trench_height_rear": 8.60473362403628e-06, + "width_overlap_front_left": 2e-06, + "width_overlap_front_right": 2e-06, + "width_overlap_rear_left": 2e-06, + "width_overlap_rear_right": 2e-06, + }, + "rough_milling_electron": { + "step_name": "Rough Milling - Electron Image", + "recipe_name": "Milling", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "site_location_type": "Chunk", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 1.25e-11, + }, + "reference_redefinition_2": { + "step_name": "Reference Redefinition 2", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 105.1459769, + "site_location_type": "Chunk", + }, + "medium_milling": { + "step_name": "Medium Milling", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 298.1471377, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 5e-10, + "depth_correction": 2.0, + "lamella_offset": 1.5e-06, + "width_overlap_front_left": 1.5e-06, + "width_overlap_front_right": 1.5e-06, + "width_overlap_rear_left": 1.5e-06, + "width_overlap_rear_right": 1.5e-06, + }, + "medium_milling_electron": { + "step_name": "Medium Milling - Electron Image", + "recipe_name": "Milling", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "site_location_type": "Chunk", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 1.25e-11, + }, + "fine_milling": { + "step_name": "Fine Milling", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 397.414388, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 3e-10, + "depth_correction": 2.0, + "lamella_offset": 1e-06, + "width_overlap_front_left": 1e-06, + "width_overlap_front_right": 1e-06, + "width_overlap_rear_left": 1e-06, + "width_overlap_rear_right": 1e-06, + }, + "fine_milling_electron": { + "step_name": "Fine Milling - Electron Image", + "recipe_name": "Milling", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "site_location_type": "Chunk", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 1.25e-11, + }, + "finer_milling": { + "step_name": "Finer Milling", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 887.9686893, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 1e-10, + "depth_correction": 2.0, + "lamella_offset": 4.0000000000000003e-07, + "width_overlap_front_left": 5.000000000000001e-07, + "width_overlap_front_right": 5.000000000000001e-07, + "width_overlap_rear_left": 5.000000000000001e-07, + "width_overlap_rear_right": 5.000000000000001e-07, + }, + "finer_milling_electron": { + "step_name": "Finer Milling - Electron Image", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 59.2296577, + "site_location_type": "Chunk", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 2.5e-11, + }, + "delay_2": { + "step_name": "Delay", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + }, + "polishing_1": { + "step_name": "Polishing 1", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "Finished", + "execution_time": 680.5567315, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 5e-11, + "depth_correction": 2.0, + "lamella_offset": 2.5000000000000004e-07, + "width_overlap_front_left": 0.0, + "width_overlap_front_right": 0.0, + "width_overlap_rear_left": 0.0, + "width_overlap_rear_right": 0.0, + }, + "polishing_1_electron": { + "step_name": "Polishing 1 - Electron Image", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "site_location_type": "Thinning", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 1.25e-11, + }, + "polishing_2": { + "step_name": "Polishing 2", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "Finished", + "execution_time": 1.170488927, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 3e-11, + "depth_correction": 1.5, + "lamella_offset": 0.0, + "width_overlap_front_left": 0.0, + "width_overlap_front_right": 0.0, + "width_overlap_rear_left": 0.0, + "width_overlap_rear_right": 0.0, + }, + "polishing_2_ion": { + "step_name": "Polishing 2 - Ion Image", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "site_location_type": "Thinning", + "beam_type": "Ion", + "voltage": 30000.0, + "current": 3e-11, + }, + "polishing_2_electron": { + "step_name": "Polishing 2 - Electron Image", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "Finished", + "execution_time": 56.9180832, + "site_location_type": "Thinning", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 2.5e-11, + }, +} +stage_info = { + "preparation_site": { + "x": -0.0031091224743875403, + "y": 0.00420867925495798, + "z": 0.0323644854106331, + "rotation": 285.003247202109, + "tilt_alpha": 25.9996646026832, + "slot_number": 1, + }, + "chunk_site": { + "x": -0.0030037500000000003, + "y": 0.004293, + "z": 0.032350405092592606, + "rotation": 285.003247202109, + "tilt_alpha": -0.000134158926728586, + "slot_number": 1, + }, + "thinning_site": { + "x": -0.0030037500000000003, + "y": 0.004293, + "z": 0.032350405092592606, + "rotation": 285.003247202109, + "tilt_alpha": -0.000134158926728586, + "slot_number": 1, + }, + "chunk_coincidence_params": { + "x": -0.0030048260286678298, + "y": 0.004308828126160981, + "z": 0.0323400707790533, + "rotation": 285.003247202109, + "tilt_alpha": -0.000134158926728586, + "slot_number": 1, + }, + "thinning_params": { + "x": -0.0030037500000000003, + "y": 0.004293, + "z": 0.032350405092592606, + "rotation": 285.003247202109, + "tilt_alpha": -0.000134158926728586, + "slot_number": 1, + }, +} site_info = { "project_name": visit_name, "site_name": "Lamella", "site_number": 1, - "stage_info": { - "preparation_site": { - "x": -0.0031091224743875403, - "y": 0.00420867925495798, - "z": 0.0323644854106331, - "rotation": 285.003247202109, - "tilt_alpha": 25.9996646026832, - "slot_number": 1, - }, - "chunk_site": { - "x": -0.0030037500000000003, - "y": 0.004293, - "z": 0.032350405092592606, - "rotation": 285.003247202109, - "tilt_alpha": -0.000134158926728586, - "slot_number": 1, - }, - "thinning_site": { - "x": -0.0030037500000000003, - "y": 0.004293, - "z": 0.032350405092592606, - "rotation": 285.003247202109, - "tilt_alpha": -0.000134158926728586, - "slot_number": 1, - }, - "chunk_coincidence_params": { - "x": -0.0030048260286678298, - "y": 0.004308828126160981, - "z": 0.0323400707790533, - "rotation": 285.003247202109, - "tilt_alpha": -0.000134158926728586, - "slot_number": 1, - }, - "thinning_params": { - "x": -0.0030037500000000003, - "y": 0.004293, - "z": 0.032350405092592606, - "rotation": 285.003247202109, - "tilt_alpha": -0.000134158926728586, - "slot_number": 1, - }, - }, - "steps": { - "eucentric_tilt": { - "step_name": "Eucentric Tilt", - "recipe_name": "Preparation", - "is_enabled": False, - "status": "Finished", - "execution_time": 328.6657458, - }, - "artificial_features": { - "step_name": "Artificial Features", - "recipe_name": "Preparation", - "is_enabled": False, - "status": "None", - "execution_time": 0.0, - "beam_type": "Ion", - "voltage": 30000.0, - "current": 5e-10, - }, - "milling_angle": { - "step_name": "Milling Angle", - "recipe_name": "Preparation", - "is_enabled": False, - "status": "Finished", - "execution_time": 208.4942922, - "milling_angle": 12.0, - }, - "image_acquisition": { - "step_name": "Image Acquisition", - "recipe_name": "Preparation", - "is_enabled": False, - "status": "Finished", - "execution_time": 19.7872887, - "site_location_type": "Chunk", - }, - "lamella_placement": { - "step_name": "Lamella Placement", - "recipe_name": "Preparation", - "is_enabled": False, - "status": "None", - "execution_time": 0.0, - }, - "delay_1": { - "step_name": "Delay", - "recipe_name": "Milling", - "is_enabled": False, - "status": "None", - "execution_time": 0.0, - }, - "reference_definition": { - "step_name": "Reference Definition", - "recipe_name": "Milling", - "is_enabled": False, - "status": "Finished", - "execution_time": 75.6320442, - "site_location_type": "Chunk", - }, - "reference_definition_electron": { - "step_name": "Electron Reference Definition", - "recipe_name": "Milling", - "is_enabled": False, - "status": "Finished", - "execution_time": 131.7462301, - }, - "stress_relief_cuts": { - "step_name": "Stress Relief Cuts", - "recipe_name": "Milling", - "is_enabled": False, - "status": "None", - "execution_time": 0.0, - "beam_type": "Ion", - "voltage": 30000.0, - "current": 1e-09, - "depth_correction": 3.0, - }, - "reference_redefinition_1": { - "step_name": "Reference Redefinition 1", - "recipe_name": "Milling", - "is_enabled": False, - "status": "None", - "execution_time": 0.0, - "site_location_type": "Chunk", - }, - "rough_milling": { - "step_name": "Rough Milling", - "recipe_name": "Milling", - "is_enabled": False, - "status": "Finished", - "execution_time": 1.5929074719, - "beam_type": "Ion", - "voltage": 30000.0, - "current": 1e-09, - "depth_correction": 3.0, - "lamella_offset": 2e-06, - "trench_height_front": 5.6338138462765e-06, - "trench_height_rear": 8.60473362403628e-06, - "width_overlap_front_left": 2e-06, - "width_overlap_front_right": 2e-06, - "width_overlap_rear_left": 2e-06, - "width_overlap_rear_right": 2e-06, - }, - "rough_milling_electron": { - "step_name": "Rough Milling - Electron Image", - "recipe_name": "Milling", - "is_enabled": False, - "status": "None", - "execution_time": 0.0, - "site_location_type": "Chunk", - "beam_type": "Electron", - "voltage": 2000.0, - "current": 1.25e-11, - }, - "reference_redefinition_2": { - "step_name": "Reference Redefinition 2", - "recipe_name": "Milling", - "is_enabled": False, - "status": "Finished", - "execution_time": 105.1459769, - "site_location_type": "Chunk", - }, - "medium_milling": { - "step_name": "Medium Milling", - "recipe_name": "Milling", - "is_enabled": False, - "status": "Finished", - "execution_time": 298.1471377, - "beam_type": "Ion", - "voltage": 30000.0, - "current": 5e-10, - "depth_correction": 2.0, - "lamella_offset": 1.5e-06, - "width_overlap_front_left": 1.5e-06, - "width_overlap_front_right": 1.5e-06, - "width_overlap_rear_left": 1.5e-06, - "width_overlap_rear_right": 1.5e-06, - }, - "medium_milling_electron": { - "step_name": "Medium Milling - Electron Image", - "recipe_name": "Milling", - "is_enabled": False, - "status": "None", - "execution_time": 0.0, - "site_location_type": "Chunk", - "beam_type": "Electron", - "voltage": 2000.0, - "current": 1.25e-11, - }, - "fine_milling": { - "step_name": "Fine Milling", - "recipe_name": "Milling", - "is_enabled": False, - "status": "Finished", - "execution_time": 397.414388, - "beam_type": "Ion", - "voltage": 30000.0, - "current": 3e-10, - "depth_correction": 2.0, - "lamella_offset": 1e-06, - "width_overlap_front_left": 1e-06, - "width_overlap_front_right": 1e-06, - "width_overlap_rear_left": 1e-06, - "width_overlap_rear_right": 1e-06, - }, - "fine_milling_electron": { - "step_name": "Fine Milling - Electron Image", - "recipe_name": "Milling", - "is_enabled": False, - "status": "None", - "execution_time": 0.0, - "site_location_type": "Chunk", - "beam_type": "Electron", - "voltage": 2000.0, - "current": 1.25e-11, - }, - "finer_milling": { - "step_name": "Finer Milling", - "recipe_name": "Milling", - "is_enabled": False, - "status": "Finished", - "execution_time": 887.9686893, - "beam_type": "Ion", - "voltage": 30000.0, - "current": 1e-10, - "depth_correction": 2.0, - "lamella_offset": 4.0000000000000003e-07, - "width_overlap_front_left": 5.000000000000001e-07, - "width_overlap_front_right": 5.000000000000001e-07, - "width_overlap_rear_left": 5.000000000000001e-07, - "width_overlap_rear_right": 5.000000000000001e-07, - }, - "finer_milling_electron": { - "step_name": "Finer Milling - Electron Image", - "recipe_name": "Milling", - "is_enabled": False, - "status": "Finished", - "execution_time": 59.2296577, - "site_location_type": "Chunk", - "beam_type": "Electron", - "voltage": 2000.0, - "current": 2.5e-11, - }, - "delay_2": { - "step_name": "Delay", - "recipe_name": "Thinning", - "is_enabled": False, - "status": "None", - "execution_time": 0.0, - }, - "polishing_1": { - "step_name": "Polishing 1", - "recipe_name": "Thinning", - "is_enabled": False, - "status": "Finished", - "execution_time": 680.5567315, - "beam_type": "Ion", - "voltage": 30000.0, - "current": 5e-11, - "depth_correction": 2.0, - "lamella_offset": 2.5000000000000004e-07, - "width_overlap_front_left": 0.0, - "width_overlap_front_right": 0.0, - "width_overlap_rear_left": 0.0, - "width_overlap_rear_right": 0.0, - }, - "polishing_1_electron": { - "step_name": "Polishing 1 - Electron Image", - "recipe_name": "Thinning", - "is_enabled": False, - "status": "None", - "execution_time": 0.0, - "site_location_type": "Thinning", - "beam_type": "Electron", - "voltage": 2000.0, - "current": 1.25e-11, - }, - "polishing_2": { - "step_name": "Polishing 2", - "recipe_name": "Thinning", - "is_enabled": False, - "status": "Finished", - "execution_time": 1.170488927, - "beam_type": "Ion", - "voltage": 30000.0, - "current": 3e-11, - "depth_correction": 1.5, - "lamella_offset": 0.0, - "width_overlap_front_left": 0.0, - "width_overlap_front_right": 0.0, - "width_overlap_rear_left": 0.0, - "width_overlap_rear_right": 0.0, - }, - "polishing_2_ion": { - "step_name": "Polishing 2 - Ion Image", - "recipe_name": "Thinning", - "is_enabled": False, - "status": "None", - "execution_time": 0.0, - "site_location_type": "Thinning", - "beam_type": "Ion", - "voltage": 30000.0, - "current": 3e-11, - }, - "polishing_2_electron": { - "step_name": "Polishing 2 - Electron Image", - "recipe_name": "Thinning", - "is_enabled": False, - "status": "Finished", - "execution_time": 56.9180832, - "site_location_type": "Thinning", - "beam_type": "Electron", - "voltage": 2000.0, - "current": 2.5e-11, - }, - }, + "stage_info": stage_info, + "steps": milling_steps, } # Construct the RabbitMQ message received message = { @@ -444,8 +446,6 @@ def test_run_with_db( assert len(gs_ispyb_rows) >= 1 gs_ispyb = gs_ispyb_rows[0] - steps = cast(dict[str, Any], site_info["steps"]) - # There should be one MillingStep entry in ISPyB for each step in the message milling_step_ispyb_rows = ( ispyb_db_session.execute( @@ -456,7 +456,7 @@ def test_run_with_db( .scalars() .all() ) - assert len(milling_step_ispyb_rows) == len(steps.keys()) + assert len(milling_step_ispyb_rows) == len(milling_steps) # There should be the same thing in Murfey milling_step_murfey_rows = murfey_db_session.exec( @@ -464,4 +464,76 @@ def test_run_with_db( MurfeyDB.MillingStep.grid_square_id == gs_ispyb.gridSquareId ) ).all() - assert len(milling_step_murfey_rows) == len(steps.keys()) + assert len(milling_step_murfey_rows) == len(milling_steps) + + +@pytest.mark.parametrize( + "test_params", + ( # TransportManager | Project name | Site number | Site name | Stage info | Preparation site | Slot number + (False, False, False, False, False, False, False), + (True, False, False, False, False, False, False), + (True, True, False, False, False, False, False), + (True, True, True, False, False, False, False), + (True, True, True, True, False, False, False), + (True, True, True, True, True, False, False), + (True, True, True, True, True, True, False), + (True, True, True, True, True, True, True), + ), +) +def test_run_fails( + mocker: MockerFixture, + test_params: tuple[bool, bool, bool, bool, bool, bool, bool], +): + # Unpack test params + ( + has_transport_object, + has_project_name, + has_site_number, + has_site_name, + has_stage_info, + has_preparation_site, + has_slot_number, + ) = test_params + + # Mock TransportManager accordingly + mock_transport_object = MagicMock() if has_transport_object else None + mocker.patch( + "murfey.workflows.fib.register_milling_progress._transport_object", + mock_transport_object, + ) + + # Mock the Murfey database + mock_murfey_db = MagicMock() + + # Mock the site info dictionary accordingly + project_name = visit_name if has_project_name else None + site_name = "Lamella" if has_site_name else None + site_num = 1 if has_site_number else None + message = { + "register": "fib.register_milling_progress", + "session_id": session_id, + "site_info": { + "project_name": project_name, + "site_name": site_name, + "site_number": site_num, + "stage_info": { + "preparation_site": { + "x": -0.0031091224743875403 if has_slot_number else None, + "y": 0.00420867925495798, + "z": 0.0323644854106331, + "rotation": 285.003247202109, + "tilt_alpha": 25.9996646026832, + "slot_number": 1, + } + if has_preparation_site + else None, + } + if has_stage_info + else None, + "steps": None, + }, + } + + # Run the function and check that the correct message is returned + result = run(message, mock_murfey_db) + assert result == {"success": False, "requeue": False}