diff --git a/src/murfey/workflows/fib/register_atlas.py b/src/murfey/workflows/fib/register_atlas.py index b9acc20e9..a3cb0ef20 100644 --- a/src/murfey/workflows/fib/register_atlas.py +++ b/src/murfey/workflows/fib/register_atlas.py @@ -251,11 +251,13 @@ def run( fib_info.session_id, metadata, murfey_db ) logger.info( - f"Registered FIB atlas image {fib_info.atlas_file} for slot {metadata.slot_number} in Murfey database" + f"Registered FIB atlas image {fib_info.atlas_file} " + f"for slot {metadata.slot_number} in Murfey database" ) except Exception: logger.error( - f"Error registering FIB atlas image {fib_info.atlas_file} in Murfey database", + "Error registering FIB atlas image " + f"{fib_info.atlas_file} in Murfey database", exc_info=True, ) return {"success": False, "requeue": False} @@ -273,7 +275,7 @@ def run( except Exception: # Log error but allow workflow to proceed logger.error( - "Exception encountered when registering data collection group for FIB workflow " + "Error registering data collection group for FIB workflow " f"for {metadata.site_name!r}: \n" f"{traceback.format_exc()}" ) diff --git a/src/murfey/workflows/fib/register_lamella_evaluation_image.py b/src/murfey/workflows/fib/register_lamella_evaluation_image.py index 5226609d1..050b11d51 100644 --- a/src/murfey/workflows/fib/register_lamella_evaluation_image.py +++ b/src/murfey/workflows/fib/register_lamella_evaluation_image.py @@ -1,19 +1,86 @@ import json import logging +import re +from datetime import datetime from pathlib import Path from typing import Any, cast from pydantic import BaseModel -from sqlmodel import Session, select +from sqlmodel import Session as SQLModelSession, select import murfey.util.db as MurfeyDB from murfey.util.config import get_machine_config from murfey.util.models import FIBImageMetadata -from murfey.workflows.fib.shared import parse_image_metadata +from murfey.workflows.fib.shared import ( + parse_image_metadata, + populate_fib_imaging_site_entry, +) logger = logging.getLogger(__name__) +# The timestamp in the lamella evaluation image follows the pattern +# yyyy-mm-dd-HH-MM-SS +# E.g. +# 2026-03-09-18-24-51_drift_corrected_image_Finer Milling - Electron Image.png +# 2026-03-10-16-06-25_drift_corrected_image_Polishing 2 - Electron Image.png +# This can be searched for using regex +# (? Character prior to pattern CANNOT be a digit +# (?!\d) --> Character after pattern CANNOT be a digit +pattern = re.compile(r"(?= current_timestamp: + fib_imaging_site = populate_fib_imaging_site_entry( + fib_imaging_site, metadata + ) + murfey_db.add(fib_imaging_site) + murfey_db.commit() + return fib_imaging_site + + class FIBLamellaImageInfo(BaseModel): session_id: int lamella_image_file: Path @@ -21,7 +88,7 @@ class FIBLamellaImageInfo(BaseModel): def run( message: dict[str, Any], - murfey_db: Session, + murfey_db: SQLModelSession, ): # Outer try-finally block to ensure the database connection is closed logger.info( @@ -77,6 +144,21 @@ def run( ) return {"success": False, "requeue": False} + try: + # Register imaging site to Murfey, or update existing one + _ = _register_fib_imaging_site(fib_info.session_id, metadata, murfey_db) + logger.info( + f"Registered lamella evaluation image {fib_info.lamella_image_file} " + f"for slot {metadata.slot_number} in Murfey database" + ) + except Exception: + logger.error( + "Error registering lamella evaluation image " + f"{fib_info.lamella_image_file} in Murfey database", + exc_info=True, + ) + return {"success": False, "requeue": False} + return {"success": True} finally: murfey_db.close() diff --git a/tests/workflows/fib/test_register_lamella_evaluation_image.py b/tests/workflows/fib/test_register_lamella_evaluation_image.py index 5018110e6..1f984315a 100644 --- a/tests/workflows/fib/test_register_lamella_evaluation_image.py +++ b/tests/workflows/fib/test_register_lamella_evaluation_image.py @@ -4,12 +4,17 @@ import pytest from pytest_mock import MockerFixture +from sqlmodel import Session as SQLModelSession, select +import murfey.util.db as MurfeyDB from murfey.util.config import MachineConfig from murfey.workflows.fib.register_lamella_evaluation_image import ( FIBImageMetadata, + FIBLamellaImageInfo, + _register_fib_imaging_site, run, ) +from murfey.workflows.fib.shared import populate_fib_imaging_site_entry from tests.conftest import ExampleVisit session_id = 10 @@ -24,12 +29,123 @@ def visit_dir(tmp_path: Path): return visit_dir +@pytest.mark.parametrize( + "test_params", + ( # Site exists | Newer site + (True, True), + (True, False), + (False, False), + ), +) +def test_register_fib_imaging_site_with_db( + test_params: tuple[bool, bool], + visit_dir: Path, + murfey_db_session: SQLModelSession, +): + # Unpack test params + has_existing_entry, newer_insert = test_params + + # Register a Session for this test + murfey_session = MurfeyDB.Session( + id=session_id, + visit=visit_name, + name=visit_name, + instrument_name=instrument_name, + started=True, + ) + murfey_db_session.add(murfey_session) + murfey_db_session.commit() + + # Lamella directory + lamella_dir = ( + visit_dir + / "autotem" + / visit_name + / "Sites" + / "Lamella" + / "LamellaEvaluationImages" + ) + # Standard metadata to use + metadata_dict = { + "voltage": 2000, + "shift_x": 0, + "shift_y": 0, + "len_x": 0.003072, + "len_y": 0.002048, + "pos_x": -0.003, + "pos_y": 0.003, + "pos_z": 0.01, + "rotation": 1.833, + "slot_number": 2, + "tilt_alpha": 0, + "tilt_beta": 0, + "pixels_x": 3072, + "pixels_y": 2048, + "pixel_size_x": 1e-6, + "pixel_size_y": 1e-6, + } + + # Create older existing entry + if has_existing_entry: + existing_timestamp = f"2026-04-16-02-39-{30 if newer_insert else 50}" + existing_file = ( + lamella_dir + / f"{existing_timestamp}_drift_corrected_image_Polishing 2 - Electron Image.png" + ) + existing_metadata = FIBImageMetadata( + visit_name=visit_name, + file=existing_file, + **metadata_dict, + ) + existing_site = MurfeyDB.ImagingSite( + session_id=session_id, + site_name=existing_metadata.site_name, + data_type="grid_square", + ) + existing_site = populate_fib_imaging_site_entry( + existing_site, existing_metadata + ) + murfey_db_session.add(existing_site) + murfey_db_session.commit() + + # Create the test image file to register + file = ( + lamella_dir + / "2026-04-16-02-39-40_drift_corrected_image_Polishing 2 - Electron Image.png" + ) + metadata = FIBImageMetadata( + visit_name=visit_name, + file=file, + **metadata_dict, + ) + + # Run the function and check that results are as expected + _register_fib_imaging_site( + session_id=session_id, + metadata=metadata, + murfey_db=murfey_db_session, + ) + + # Only one entry should exist + found_sites = murfey_db_session.exec(select(MurfeyDB.ImagingSite)).all() + assert len(found_sites) == 1 + + # Key parameters should be populated + registered_site = found_sites[0] + assert registered_site.session_id == session_id + assert registered_site.data_type == "grid_square" + assert registered_site.site_name == metadata.site_name + if has_existing_entry and not newer_insert: + assert registered_site.image_path != str(file) + else: + assert registered_site.image_path == str(file) + + def test_run( mocker: MockerFixture, visit_dir: Path, ): # Set up parameters - project_name = "some_project" # Mock the logger mock_logger = mocker.patch( @@ -56,7 +172,7 @@ def test_run( file = ( visit_dir / "autotem" - / project_name + / visit_name / "Sites" / "Lamella" / "LamellaEvaluationImages" @@ -92,17 +208,32 @@ def test_run( return_value=metadata_dict, ) + # Mock the results of '_register_fib_image_site' + mock_register_imaging_site = mocker.patch( + "murfey.workflows.fib.register_lamella_evaluation_image._register_fib_imaging_site", + return_value=MagicMock(), + ) + # Construct the message to pass to the function message = { "register": "fib.register_lamella_evaluation_image", "session_id": session_id, "lamella_image_file": str(file), } + fib_info = FIBLamellaImageInfo(**message) # Run function and check that expected calls were made result = run(message, mock_murfey_db) - mock_logger.info.assert_called_with( + + # Metadata should have been extracted and logged + mock_logger.info.assert_any_call( "Extracted the following metadata from the image:\n" f"{json.dumps(metadata.model_dump(), indent=2, default=str)}" ) + # Imaging site registration function should have been called + mock_register_imaging_site.assert_called_once_with( + fib_info.session_id, + metadata, + mock_murfey_db, + ) assert result["success"]