diff --git a/src/murfey/client/contexts/fib.py b/src/murfey/client/contexts/fib.py index 874fc6b94..0b06370d2 100644 --- a/src/murfey/client/contexts/fib.py +++ b/src/murfey/client/contexts/fib.py @@ -1,7 +1,6 @@ from __future__ import annotations import logging -import re import threading import xml.etree.ElementTree as ET from dataclasses import dataclass, field @@ -11,6 +10,7 @@ from murfey.client.context import Context from murfey.client.instance_environment import MurfeyInstanceEnvironment from murfey.util.client import capture_post +from murfey.util.fib import number_from_name from murfey.util.models import ( LamellaSiteInfo, MillingStepInfo, @@ -24,22 +24,6 @@ lock = threading.Lock() -def _number_from_name(name: str) -> int: - """ - In the AutoTEM and Maps workflows for the FIB, the sites and images are - auto-incremented with parenthesised numbers (e.g. "Lamella (2)"), with - the first site/image typically not having a number. - - This function extracts the number from the file name, and returns 1 if - no such number is found. - """ - return ( - int(match.group(1)) - if (match := re.search(r"^[\w\s]+\((\d+)\)$", name)) is not None - else 1 - ) - - T = TypeVar("T") @@ -416,7 +400,7 @@ def _parse_autotem_metadata(self, file: Path): if (site_name := _parse_xml_text(site, "Name", str)) is None: logger.warning("Current site doesn't have a name") continue - site_num = _number_from_name(site_name) + site_num = number_from_name(site_name) site_info = LamellaSiteInfo( project_name=project_name, site_name=site_name, @@ -555,7 +539,7 @@ def _make_drift_correction_gif( parts = file.parts try: lamella_name = parts[parts.index("Sites") + 1] - lamella_number = _number_from_name(lamella_name) + lamella_number = number_from_name(lamella_name) except Exception: logger.warning( f"Could not extract metadata from file {file}", exc_info=True diff --git a/src/murfey/server/api/workflow_fib.py b/src/murfey/server/api/workflow_fib.py index f40629fa0..577a11a4d 100644 --- a/src/murfey/server/api/workflow_fib.py +++ b/src/murfey/server/api/workflow_fib.py @@ -1,7 +1,6 @@ import json import logging import os -from importlib.metadata import entry_points from pathlib import Path import numpy as np @@ -11,6 +10,7 @@ from sqlmodel import Session, select import murfey.util.db as MurfeyDB +from murfey.server import _transport_object from murfey.server.api.auth import validate_instrument_token from murfey.server.murfey_db import murfey_db from murfey.util import sanitise_path @@ -26,30 +26,25 @@ ) -class FIBAtlasInfo(BaseModel): - file: Path | None = None +class FIBAtlasFile(BaseModel): + file: Path @router.post("/sessions/{session_id}/register_atlas") def register_fib_atlas( session_id: int, - fib_atlas_info: FIBAtlasInfo, - db: Session = murfey_db, + fib_atlas: FIBAtlasFile, ): - # See if the relevant workflow is available - if not ( - workflow_search := list( - entry_points(group="murfey.workflows", name="fib.register_atlas") - ) - ): - raise RuntimeError("Unable to find Murfey workflow to register FIB atlas") - workflow = workflow_search[0] - - # Run the workflow - workflow.load()( - session_id=session_id, - file=fib_atlas_info.file, - murfey_db=db, + if _transport_object is None: + logger.error("No Transport Manager object was set up") + return None + _transport_object.send( + _transport_object.feedback_queue, + { + "register": "fib.register_atlas", + "session_id": session_id, + "atlas_file": str(fib_atlas.file), + }, ) diff --git a/src/murfey/util/db.py b/src/murfey/util/db.py index 5ba6691d6..eb92f48fa 100644 --- a/src/murfey/util/db.py +++ b/src/murfey/util/db.py @@ -10,7 +10,9 @@ from sqlmodel import Enum, Field, Relationship, SQLModel, create_engine """ +======================================================================================= GENERAL +======================================================================================= """ @@ -173,7 +175,9 @@ class ImagingSite(SQLModel, table=True): # type: ignore """ +======================================================================================= TEM SESSION AND PROCESSING WORKFLOW +======================================================================================= """ @@ -1072,7 +1076,9 @@ class CryoemInitialModel(SQLModel, table=True): # type: ignore """ +======================================================================================= FUNCTIONS +======================================================================================= """ diff --git a/src/murfey/util/fib.py b/src/murfey/util/fib.py new file mode 100644 index 000000000..7a5efdadb --- /dev/null +++ b/src/murfey/util/fib.py @@ -0,0 +1,21 @@ +""" +General functinos specific to the FIB workflow +""" + +import re + + +def number_from_name(name: str) -> int: + """ + In the AutoTEM and Maps workflows for the FIB, the sites and images are + auto-incremented with parenthesised numbers (e.g. "Lamella (2)"), with + the first site/image typically not having a number. + + This function extracts the number from the file name, and returns 1 if + no such number is found. + """ + return ( + int(match.group(1)) + if (match := re.search(r"^[\w\s]+\((\d+)\)$", name)) is not None + else 1 + ) diff --git a/src/murfey/workflows/fib/register_atlas.py b/src/murfey/workflows/fib/register_atlas.py index ebed1a198..03df5dc19 100644 --- a/src/murfey/workflows/fib/register_atlas.py +++ b/src/murfey/workflows/fib/register_atlas.py @@ -1,7 +1,10 @@ import logging +import traceback import xml.etree.ElementTree as ET from functools import cached_property +from importlib.metadata import entry_points from pathlib import Path +from typing import Any import numpy as np import PIL.Image @@ -9,6 +12,7 @@ from sqlmodel import Session, select import murfey.util.db as MurfeyDB +from murfey.util.fib import number_from_name logger = logging.getLogger("murfey.workflows.fib.register_atlas") @@ -21,6 +25,7 @@ class FIBAtlasMetadata(BaseModel): visit_name: str file: Path + thumbnail_path: Path | None = None # Acceleration voltage voltage: float # Beam shifts @@ -75,16 +80,24 @@ def slot_number(self) -> int: # mypy doesn't support decorators on @property @computed_field # type: ignore @cached_property - def site_name(self) -> str: + def project_name(self) -> str: """ - Create a site name for the current image based on the project name - and its slot number. This assumes a specific folder structure of - {visit_name}/maps/{project_name} + Extract the project name from the file path. This assumes a specific + folder structure of '{visit_name}/maps/{project_name}'. """ path_parts = self.file.parts visit_idx = path_parts.index(self.visit_name) - project_name = path_parts[visit_idx + 2] # {visit}/maps/{project_name} - return f"{project_name}--slot_{self.slot_number}" + return path_parts[visit_idx + 2] # {visit}/maps/{project_name} + + # mypy doesn't support decorators on @property + @computed_field # type: ignore + @cached_property + def site_name(self) -> str: + """ + Create a site name for the current image based on the project name + and its slot number. + """ + return f"{self.project_name}--slot_{self.slot_number}" def _parse_metadata(file: Path, visit_name: str): @@ -142,6 +155,31 @@ def _parse_metadata(file: Path, visit_name: str): ) +def _make_thumbnail(file: Path, metadata: FIBAtlasMetadata, visit_name: str): + img = PIL.Image.open(file) + img.thumbnail((512, 512)) + + # Find visit directory path + visit_idx = file.parts.index(visit_name) + visit_dir = list(reversed(file.parents))[visit_idx] + + # Construct path to thumbnail + processed_dir = visit_dir / "processed" + image_number = number_from_name(file.stem) + save_path = ( + processed_dir + / metadata.project_name + / f"grid_{metadata.slot_number}" + / "atlas" + / f"atlas_{str(image_number).zfill(2)}.png" + ) + save_path.parent.mkdir(parents=True, exist_ok=True) + + # Save the thumbnail + img.save(save_path) + return save_path + + def _register_fib_imaging_site( session_id: int, metadata: FIBAtlasMetadata, @@ -150,75 +188,244 @@ def _register_fib_imaging_site( """ Register FIB atlas in Murfey database or update existing entry. """ - # Create new entry if one doesn't already exist - if not ( + + def _update_entry( + imaging_site: MurfeyDB.ImagingSite, + metadata: FIBAtlasMetadata, + ): + imaging_site.image_path = str(metadata.file) + imaging_site.pos_x = metadata.pos_x + imaging_site.pos_y = metadata.pos_y + imaging_site.pos_z = metadata.pos_z + imaging_site.rotation = float(np.rad2deg(metadata.rotation)) + imaging_site.tilt_alpha = float(np.rad2deg(metadata.tilt_alpha)) + imaging_site.tilt_beta = float(np.rad2deg(metadata.tilt_beta)) + imaging_site.len_x = metadata.len_x + imaging_site.len_y = metadata.len_y + imaging_site.image_pixels_x = metadata.pixels_x + imaging_site.image_pixels_y = metadata.pixels_y + imaging_site.image_pixel_size = metadata.pixel_size + + if metadata.thumbnail_path is not None: + scale = 512 / (max(metadata.pixels_x, metadata.pixels_y) or 1) + imaging_site.thumbnail_path = str(metadata.thumbnail_path) + imaging_site.thumbnail_pixels_x = int(round(metadata.pixels_x * scale)) or 1 + imaging_site.thumbnail_pixels_y = int(round(metadata.pixels_y * scale)) or 1 + imaging_site.thumbnail_pixel_size = metadata.pixel_size / scale + + return imaging_site + + if ( fib_imaging_site := murfey_db.exec( select(MurfeyDB.ImagingSite) .where(MurfeyDB.ImagingSite.session_id == session_id) - .where(MurfeyDB.ImagingSite.image_path == str(metadata.file)) + .where(MurfeyDB.ImagingSite.site_name == metadata.site_name) + .where(MurfeyDB.ImagingSite.data_type == "atlas") ).one_or_none() - ): + ) is None: + # Create new entry if one doesn't already exist fib_imaging_site = MurfeyDB.ImagingSite( session_id=session_id, + site_name=metadata.site_name, image_path=str(metadata.file), data_type="atlas", ) - # Add/update entries - fib_imaging_site.site_name = metadata.site_name - fib_imaging_site.pos_x = metadata.pos_x - fib_imaging_site.pos_y = metadata.pos_y - fib_imaging_site.pos_z = metadata.pos_z - fib_imaging_site.rotation = float(np.rad2deg(metadata.rotation)) - fib_imaging_site.tilt_alpha = float(np.rad2deg(metadata.tilt_alpha)) - fib_imaging_site.tilt_beta = float(np.rad2deg(metadata.tilt_beta)) - fib_imaging_site.len_x = metadata.len_x - fib_imaging_site.len_y = metadata.len_y - fib_imaging_site.image_pixels_x = metadata.pixels_x - fib_imaging_site.image_pixels_y = metadata.pixels_y - fib_imaging_site.image_pixel_size = metadata.pixel_size + fib_imaging_site = _update_entry(fib_imaging_site, metadata) + else: + # Check if the entry is new or newer than the current stored one + incoming_number = number_from_name(metadata.file.stem) + # Handle empty string + if not fib_imaging_site.image_path: + current_number = 0 + # Read 'maps' atlases in one way + elif "maps" in (curr_path := Path(fib_imaging_site.image_path)).parts: + current_number = number_from_name(curr_path.stem) + else: + current_number = 0 + # Update if incoming one is newer + if incoming_number >= current_number: + fib_imaging_site = _update_entry(fib_imaging_site, metadata) murfey_db.add(fib_imaging_site) murfey_db.commit() + return fib_imaging_site -def run( + +def _register_dcg_and_atlas( session_id: int, - file: Path, + instrument_name: str, + visit_name: str, + imaging_site: MurfeyDB.ImagingSite, + metadata: FIBAtlasMetadata, + murfey_db: Session, +): + 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 using thumbnail values if they are provided + if ( + imaging_site.thumbnail_path is not None + and imaging_site.thumbnail_pixel_size is not None + ): + atlas_name: str | None = imaging_site.thumbnail_path + atlas_pixel_size: float | None = imaging_site.thumbnail_pixel_size + else: + atlas_name = imaging_site.image_path + atlas_pixel_size = imaging_site.image_pixel_size + + if dcg_search := murfey_db.exec( + select(MurfeyDB.DataCollectionGroup) + .where(MurfeyDB.DataCollectionGroup.session_id == session_id) + .where(MurfeyDB.DataCollectionGroup.tag == imaging_site.site_name) + ).all(): + dcg_entry = dcg_search[0] + atlas_message = { + "session_id": session_id, + "dcgid": dcg_entry.id, + "atlas_id": dcg_entry.atlas_id, + "atlas": atlas_name, + "atlas_pixel_size": atlas_pixel_size, + "sample": dcg_entry.sample, + } + if entry_point_result := entry_points( + group="murfey.workflows", name="atlas_update" + ): + (workflow,) = entry_point_result + _ = workflow.load()( + message=atlas_message, + murfey_db=murfey_db, + ) + else: + logger.warning("No workflow found for 'atlas_update'") + else: + dcg_message = { + "microscope": instrument_name, + "proposal_code": proposal_code, + "proposal_number": proposal_number, + "visit_number": visit_number, + "session_id": session_id, + "tag": imaging_site.site_name, + "experiment_type_id": 46, + "atlas": atlas_name, + "atlas_pixel_size": atlas_pixel_size, + "sample": metadata.slot_number, + } + if entry_point_result := entry_points( + group="murfey.workflows", name="data_collection_group" + ): + (workflow,) = entry_point_result + # Register grid square + _ = workflow.load()( + message=dcg_message, + murfey_db=murfey_db, + ) + else: + logger.warning("No workflow found for 'data_collection_group'") + dcg_entry = murfey_db.exec( + select(MurfeyDB.DataCollectionGroup) + .where(MurfeyDB.DataCollectionGroup.session_id == session_id) + .where(MurfeyDB.DataCollectionGroup.tag == imaging_site.site_name) + ).one() + + imaging_site.dcg_id = dcg_entry.id + imaging_site.dcg_name = dcg_entry.tag + murfey_db.add(imaging_site) + murfey_db.commit() + + +class FIBAtlasRegistrationInfo(BaseModel): + register: str + session_id: int + atlas_file: Path + + +def run( + message: dict[str, Any], murfey_db: Session, ): # Outer try-finally block to ensure database connection closes try: - # Load visit information try: - session_entry = murfey_db.exec( - select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id) + # Validate incoming message + fib_info = FIBAtlasRegistrationInfo(**message) + except Exception: + logger.error("Could not validate incoming message", exc_info=True) + return {"success": False, "requeue": False} + + try: + # Load visit information + murfey_session = murfey_db.exec( + select(MurfeyDB.Session).where( + MurfeyDB.Session.id == fib_info.session_id + ) ).one() - visit_name = session_entry.visit + visit_name = murfey_session.visit except Exception: logger.error( "Exception encountered while querying Murfey database", exc_info=True ) - return False + return {"success": False, "requeue": False} - # Extract metadata from Electron Snapshot image try: - metadata = _parse_metadata(file, visit_name) + # Extract metadata from Electron Snapshot image + metadata = _parse_metadata(fib_info.atlas_file, visit_name) except Exception: - logger.error(f"Error extracting metadata from file {file}", exc_info=True) - return False + logger.error( + f"Error extracting metadata from file {fib_info.atlas_file}", + exc_info=True, + ) + return {"success": False, "requeue": False} + + try: + # Make a thumbnail of the image and update metadata accordingly + metadata.thumbnail_path = _make_thumbnail( + file=metadata.file, + metadata=metadata, + visit_name=visit_name, + ) + except Exception: + logger.warning( + f"Error creating thumbnail of file {fib_info.atlas_file}", exc_info=True + ) - # Register imaging site in Murfey, or update existing one try: - _register_fib_imaging_site(session_id, metadata, murfey_db) + # Register imaging site in Murfey, or update existing one + fib_imaging_site = _register_fib_imaging_site( + fib_info.session_id, metadata, murfey_db + ) logger.info( - f"Registered FIB atlas image {file} for slot {metadata.slot_number} in Murfey database" + f"Registered FIB atlas image {fib_info.atlas_file} for slot {metadata.slot_number} in Murfey database" ) except Exception: logger.error( - f"Error registering FIB atlas image {file} in Murfey database", + f"Error registering FIB atlas image {fib_info.atlas_file} in Murfey database", exc_info=True, ) - return False - return True + return {"success": False, "requeue": False} + + try: + # Register data collection group and atlas in ISPyB + _register_dcg_and_atlas( + session_id=fib_info.session_id, + instrument_name=murfey_session.instrument_name, + visit_name=murfey_session.visit, + imaging_site=fib_imaging_site, + metadata=metadata, + murfey_db=murfey_db, + ) + except Exception: + # Log error but allow workflow to proceed + logger.error( + "Exception encountered when registering data collection group for FIB workflow " + f"for {metadata.site_name!r}: \n" + f"{traceback.format_exc()}" + ) + return {"success": False, "requeue": True} + return {"success": True, "requeue": False} + finally: murfey_db.close() diff --git a/tests/client/contexts/test_fib.py b/tests/client/contexts/test_fib.py index f13c552f4..218a669c5 100644 --- a/tests/client/contexts/test_fib.py +++ b/tests/client/contexts/test_fib.py @@ -14,7 +14,6 @@ FIBImage, _file_transferred_to, _get_source, - _number_from_name, _parse_boolean, ) from murfey.util.models import LamellaSiteInfo @@ -392,24 +391,6 @@ def fib_maps_images(visit_dir: Path): # ------------------------------------------------------------------------------------- -@pytest.mark.parametrize( - "test_params", - ( # File name | Expected number - # AutoTEM examples - ("Lamella", 1), - ("Lamella (2)", 2), - ("Lamella (12)", 12), - # Maps examples - ("Electron Snapshot", 1), - ("Electron Snapshot (3)", 3), - ("Electron Snapshot (21)", 21), - ), -) -def test_number_from_name(test_params: tuple[str, int]): - name, number = test_params - assert _number_from_name(name) == number - - @pytest.mark.parametrize( "test_params", ( # Input | Expected output diff --git a/tests/server/api/test_workflow_fib.py b/tests/server/api/test_workflow_fib.py index 3afdb935f..28b37ae29 100644 --- a/tests/server/api/test_workflow_fib.py +++ b/tests/server/api/test_workflow_fib.py @@ -7,58 +7,64 @@ from pytest_mock import MockerFixture from murfey.server.api.workflow_fib import ( - FIBAtlasInfo, + FIBAtlasFile, FIBGIFParameters, make_gif, register_fib_atlas, ) +@pytest.mark.parametrize( + "has_transport_object", + ( + True, + False, + ), +) def test_register_fib_atlas( mocker: MockerFixture, tmp_path: Path, + has_transport_object: bool, ): - # Mock the databse instance - mock_db = MagicMock() - - # Patch out the entry point being called - mock_register_fib_atlas = mocker.patch("murfey.workflows.fib.register_atlas.run") - + # Set up the variables session_id = 1 - fib_atlas_info = FIBAtlasInfo(**{"file": str(tmp_path / "dummy")}) + fib_atlas = FIBAtlasFile(**{"file": str(tmp_path / "dummy")}) + + # Mock the logger + mock_logger = mocker.patch("murfey.server.api.workflow_fib.logger") + + # Mock the tranposrt object + if has_transport_object: + mock_transport_object = MagicMock() + mock_transport_object.feedback_queue = "dummy" + mocker.patch( + "murfey.server.api.workflow_fib._transport_object", + mock_transport_object, + ) + else: + mocker.patch( + "murfey.server.api.workflow_fib._transport_object", + None, + ) # Run the function and check that the expected calls were made register_fib_atlas( session_id=session_id, - fib_atlas_info=fib_atlas_info, - db=mock_db, - ) - mock_register_fib_atlas.assert_called_once_with( - session_id=session_id, - file=fib_atlas_info.file, - murfey_db=mock_db, + fib_atlas=fib_atlas, ) - -def test_register_fib_atlas_no_entry_point( - mocker: MockerFixture, - tmp_path: Path, -): - # Mock out entry_points to return an empty list - mocker.patch("murfey.server.api.workflow_fib.entry_points", return_value=[]) - - # Mock the databse instance - mock_db = MagicMock() - - fib_atlas_info = FIBAtlasInfo(**{"file": str(tmp_path / "dummy")}) - - # Patch out the entry point being called - with pytest.raises(RuntimeError): - register_fib_atlas( - session_id=1, - fib_atlas_info=fib_atlas_info, - db=mock_db, + # Check that the expected calls were made + if has_transport_object: + mock_transport_object.send.assert_called_with( + "dummy", + { + "register": "fib.register_atlas", + "session_id": session_id, + "atlas_file": str(fib_atlas.file), + }, ) + else: + mock_logger.error.assert_called_with("No Transport Manager object was set up") @pytest.mark.asyncio diff --git a/tests/util/test_fib_util.py b/tests/util/test_fib_util.py new file mode 100644 index 000000000..6a8d9a73b --- /dev/null +++ b/tests/util/test_fib_util.py @@ -0,0 +1,21 @@ +import pytest + +from murfey.util.fib import number_from_name + + +@pytest.mark.parametrize( + "test_params", + ( # File name | Expected number + # AutoTEM examples + ("Lamella", 1), + ("Lamella (2)", 2), + ("Lamella (12)", 12), + # Maps examples + ("Electron Snapshot", 1), + ("Electron Snapshot (3)", 3), + ("Electron Snapshot (21)", 21), + ), +) +def test_number_from_name(test_params: tuple[str, int]): + name, number = test_params + assert number_from_name(name) == number diff --git a/tests/workflows/fib/test_register_atlas.py b/tests/workflows/fib/test_register_atlas.py index 3b2600b51..455a62ad4 100644 --- a/tests/workflows/fib/test_register_atlas.py +++ b/tests/workflows/fib/test_register_atlas.py @@ -2,16 +2,22 @@ from pathlib import Path from unittest.mock import MagicMock +import ispyb.sqlalchemy as ISPyBDB +import numpy as np +import PIL.Image import pytest from pytest_mock import MockerFixture -from sqlmodel import Session, select +from sqlalchemy import select as sa_select +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_atlas import FIBAtlasMetadata, _parse_metadata, run +from tests.conftest import ExampleVisit session_id = 10 -visit_name = "cm12345-6" -instrument_name = "test_instrument" +visit_name = f"{ExampleVisit.proposal_code}{ExampleVisit.proposal_number}-{ExampleVisit.visit_number}" +instrument_name = ExampleVisit.instrument_name @pytest.fixture @@ -289,16 +295,20 @@ def test_register_fib_imaging_site(): def test_run_with_db( mocker: MockerFixture, visit_dir: Path, - murfey_db_session: Session, + murfey_db_session: SQLModelSession, + ispyb_db_session: SQLAlchemySession, + mock_ispyb_credentials, ): - test_file = ( - visit_dir / "maps/LayersData/Layer/Electron Snapshot/Electron Snapshot.tiff" + test_files = ( + visit_dir / "maps/LayersData/Layer/Electron Snapshot/Electron Snapshot.tiff", + visit_dir + / "maps/LayersData/Layer/Electron Snapshot/Electron Snapshot (2).tiff", ) # Add a test visit to the database if not ( session_entry := murfey_db_session.exec( - select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id) + sm_select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id) ).one_or_none() ): session_entry = MurfeyDB.Session(id=session_id) @@ -309,34 +319,132 @@ def test_run_with_db( murfey_db_session.add(session_entry) murfey_db_session.commit() - # Mock the metadata returned from the image file - mock_metadata = FIBAtlasMetadata( - visit_name=visit_name, - file=test_file, - voltage=2000, - shift_x=0, - shift_y=0, - len_x=0.003072, - len_y=0.002048, - pos_x=0.003, - pos_y=0.0003, - pos_z=0.01, - rotation=-1.309, - tilt_alpha=0.8, - tilt_beta=0, - pixels_x=3072, - pixels_y=2048, - pixel_size_x=1e-6, - pixel_size_y=1e-6, + # 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 data collection group + 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.register_atlas_update._transport_object", + new=TransportManager("PikaTransport"), + ) + + # Mock the metadata returned from the image file + import murfey.workflows.fib.register_atlas + + mock_metadata = [ + FIBAtlasMetadata( + visit_name=visit_name, + file=test_file, + voltage=2000, + shift_x=0, + shift_y=0, + len_x=0.003072, + len_y=0.002048, + pos_x=0.003, + pos_y=0.0003, + pos_z=0.01, + rotation=-1.309, + tilt_alpha=0.8, + tilt_beta=0, + pixels_x=3072, + pixels_y=2048, + pixel_size_x=1e-6, + pixel_size_y=1e-6, + ) + for test_file in test_files + ] + mock_parse = mocker.patch( "murfey.workflows.fib.register_atlas._parse_metadata", - return_value=mock_metadata, + side_effect=mock_metadata, + ) + spy_register = mocker.spy( + murfey.workflows.fib.register_atlas, + "_register_fib_imaging_site", + ) + + # Mock 'PIL.Image.open' and create a test image + mocker.patch( + "murfey.workflows.fib.register_atlas.PIL.Image.open", + return_value=PIL.Image.fromarray(np.ones((2048, 1152), dtype=np.uint8)), ) # Run the function and check that it's run through to completion - assert run( - session_id=session_id, - file=test_file, - murfey_db=murfey_db_session, + for test_file in test_files: + run( + message={ + "register": "fib.register_atlas", + "session_id": session_id, + "atlas_file": str(test_file), + }, + murfey_db=murfey_db_session, + ) + assert mock_parse.call_count == len(test_files) + assert spy_register.call_count == len(test_files) + + # Murfey's ImagingSite should have an entry + search_results = murfey_db_session.exec( + sm_select(MurfeyDB.ImagingSite).where( + MurfeyDB.ImagingSite.session_id == session_id + ) + ).all() + assert len(search_results) == 1 + assert search_results[0].image_path == str(mock_metadata[-1].file) + + # Murfey's DataCollectionGroup should have an entry + murfey_dcg_search = murfey_db_session.exec( + sm_select(MurfeyDB.DataCollectionGroup).where( + MurfeyDB.DataCollectionGroup.session_id == session_id + ) + ).all() + assert len(murfey_dcg_search) == 1 + + # ISPyB's DataCollectionGroup should have an entry + murfey_dcg = murfey_dcg_search[0] + ispyb_dcg_search = ( + ispyb_db_session.execute( + sa_select(ISPyBDB.DataCollectionGroup).where( + ISPyBDB.DataCollectionGroup.dataCollectionGroupId == murfey_dcg.id + ) + ) + .scalars() + .all() + ) + assert len(ispyb_dcg_search) == 1 + + # Atlas should have an entry + ispyb_dcg = ispyb_dcg_search[0] + ispyb_atlas_search = ( + ispyb_db_session.execute( + sa_select(ISPyBDB.Atlas).where( + ISPyBDB.Atlas.dataCollectionGroupId == ispyb_dcg.dataCollectionGroupId + ) + ) + .scalars() + .all() + ) + assert len(ispyb_atlas_search) == 1 + ispyb_atlas_entry = ispyb_atlas_search[0] + assert ispyb_atlas_entry.atlasImage.endswith( + f"atlas_{str(mock_metadata[-1].slot_number).zfill(2)}.png" )