test_register_preprocessing_results.py

from pathlib import Path
from typing import Any
from unittest.mock import MagicMock

import ispyb.sqlalchemy as ISPyBDB
import pytest
from pytest_mock import MockerFixture
from sqlalchemy import select as sa_select
from sqlalchemy.orm.session import Session as SQLAlchemySession
from sqlmodel import select as sm_select
from sqlmodel.orm.session import Session as SQLModelSession

import murfey.util.db as MurfeyDB
from murfey.workflows.clem.register_preprocessing_results import (
    COLOR_FLAGS_MURFEY_TO_ISPYB,
    _determine_collection_mode,
    _get_color_flags,
    _register_clem_imaging_site,
    _register_dcg_and_atlas,
    _register_grid_square,
    _snake_to_camel_case,
    run,
)
from tests.conftest import ExampleVisit, get_or_create_db_entry

visit_name = f"{ExampleVisit.proposal_code}{ExampleVisit.proposal_number}-{ExampleVisit.visit_number}"
processed_dir_name = "processed"
grid_name = "Grid_1"


@pytest.fixture
def rsync_basepath(tmp_path: Path):
    return tmp_path / "data"


def generate_preprocessing_messages(
    rsync_basepath: Path,
    session_id: int,
    colors: list[str],
):
    # Make directory to where data for current grid is stored
    visit_dir = rsync_basepath / "2020" / visit_name
    processed_dir = visit_dir / processed_dir_name
    grid_dir = processed_dir / grid_name
    grid_dir.mkdir(parents=True, exist_ok=True)

    # Construct all the datasets to be tested
    datasets: list[tuple[Path, bool, bool, tuple[int, int], float, list[float]]] = [
        (
            grid_dir / "Overview 1" / "Image 1",
            False,
            True,
            (2400, 2400),
            1e-6,
            [0.002, 0.0044, 0.002, 0.0044],
        )
    ]
    # Add on metadata for a few grid squares
    datasets.extend(
        [
            (
                grid_dir / "TileScan 1" / f"Position {n + 1}",
                True,
                False,
                (2048, 2048),
                1.6e-7,
                [0.003, 0.00332768, 0.003, 0.00332768],
            )
            for n in range(3)
        ]
    )
    datasets.extend(
        [
            (
                grid_dir / "TileScan 1" / f"Position {n + 1}_Lng_LVCC",
                True,
                False,
                (2048, 2048),
                1.6e-7,
                [0.003, 0.00332768, 0.003, 0.00332768],
            )
            for n in range(3)
        ]
    )

    messages: list[dict[str, Any]] = []
    for series_path, is_stack, is_montage, shape, pixel_size, extent in datasets:
        # Unpack items from list of dataset parameters
        series_name = (
            str(series_path.relative_to(processed_dir))
            .replace("/", "--")
            .replace(" ", "_")
        )
        metadata = series_path / "metadata" / f"{series_path.stem}.xml"
        metadata.parent.mkdir(parents=True, exist_ok=True)
        metadata.touch(exist_ok=True)
        output_files = {color: str(series_path / f"{color}.tiff") for color in colors}
        for output_file in output_files.values():
            Path(output_file).touch(exist_ok=True)
        thumbnails = {
            color: str(series_path / ".thumbnails" / f"{color}.png") for color in colors
        }
        for v in thumbnails.values():
            if not (thumbnail := Path(v)).parent.exists():
                thumbnail.parent.mkdir(parents=True)
            thumbnail.touch(exist_ok=True)
        thumbnail_size = (512, 512)

        message = {
            "session_id": session_id,
            "result": {
                "series_name": series_name,
                "number_of_members": 3,
                "is_stack": is_stack,
                "is_montage": is_montage,
                "output_files": output_files,
                "thumbnails": thumbnails,
                "thumbnail_size": thumbnail_size,
                "metadata": str(metadata),
                "parent_lif": None,
                "parent_tiffs": {},
                "pixels_x": shape[0],
                "pixels_y": shape[1],
                "units": "m",
                "pixel_size": pixel_size,
                "resolution": 1 / pixel_size,
                "extent": extent,
            },
        }
        messages.append(message)
    return messages


@pytest.mark.parametrize(
    "test_params",
    (
        (
            ["gray"],
            {
                "has_grey": True,
            },
        ),
        (
            ["gray", "red"],
            {
                "has_grey": True,
                "has_red": True,
            },
        ),
        (
            ["red", "green", "blue"],
            {
                "has_red": True,
                "has_green": True,
                "has_blue": True,
            },
        ),
        (
            ["cyan", "magenta", "yellow"],
            {
                "has_cyan": True,
                "has_magenta": True,
                "has_yellow": True,
            },
        ),
    ),
)
def test_get_color_flags(test_params: tuple[list[str], dict[str, bool]]):
    colors, positive_flags = test_params
    expected_result = dict.fromkeys(
        (
            "has_grey",
            "has_red",
            "has_green",
            "has_blue",
            "has_cyan",
            "has_magenta",
            "has_yellow",
        ),
        False,
    )
    for flag, value in positive_flags.items():
        expected_result[flag] = value
    assert _get_color_flags(colors) == expected_result


def test_register_clem_imaging_site():
    _register_clem_imaging_site


@pytest.mark.parametrize(
    "test_params",
    (
        (["gray"], "Bright Field"),
        (["gray", "blue"], "Bright Field and Fluorescent"),
        (["red", "green", "blue"], "Fluorescent"),
    ),
)
def test_determine_collection_mode(test_params: tuple[list[str], str]):
    colors, expected_result = test_params
    assert _determine_collection_mode(colors) == expected_result


@pytest.mark.parametrize(
    "test_params",
    (
        ("has_grey", "hasGrey"),
        ("has_red", "hasRed"),
        ("has_green", "hasGreen"),
        ("has_blue", "hasBlue"),
        ("has_cyan", "hasCyan"),
        ("has_magenta", "hasMagenta"),
        ("has_yellow", "hasYellow"),
    ),
)
def test_snake_to_camel_case(
    test_params: tuple[str, str],
):
    string, expected_result = test_params
    assert _snake_to_camel_case(string) == expected_result


def test_register_dcg_and_atlas():
    _register_dcg_and_atlas


def test_register_grid_square():
    _register_grid_square


@pytest.mark.parametrize(
    "test_params",
    (  # Colors
        (["gray"],),
        (["gray", "green"],),
        (["red", "green", "blue"],),
        (["cyan", "magenta", "blue"],),
    ),
)
def test_run(
    mocker: MockerFixture,
    rsync_basepath: Path,
    test_params: tuple[list[str]],
):
    # Unpack test params
    (colors,) = test_params

    # Mock the MurfeyDB connection
    mock_murfey_session_entry = MagicMock()
    mock_murfey_session_entry.instrument_name = ExampleVisit.instrument_name
    mock_murfey_session_entry.visit = visit_name
    mock_murfey_db = MagicMock()
    mock_murfey_db.exec().return_value.one.return_value = mock_murfey_session_entry

    # Mock the registration helper functions
    mock_register_clem_series = mocker.patch(
        "murfey.workflows.clem.register_preprocessing_results._register_clem_imaging_site"
    )
    mock_register_dcg_and_atlas = mocker.patch(
        "murfey.workflows.clem.register_preprocessing_results._register_dcg_and_atlas"
    )
    mock_register_grid_square = mocker.patch(
        "murfey.workflows.clem.register_preprocessing_results._register_grid_square"
    )

    # Mock the align and merge workflow call
    mock_align_and_merge_call = mocker.patch(
        "murfey.workflows.clem.register_preprocessing_results.run_align_and_merge"
    )

    preprocessing_messages = generate_preprocessing_messages(
        rsync_basepath=rsync_basepath,
        session_id=ExampleVisit.murfey_session_id,
        colors=colors,
    )
    for message in preprocessing_messages:
        result = run(
            message=message,
            murfey_db=mock_murfey_db,
        )
        assert result == {"success": True}
    assert mock_register_clem_series.call_count == len(preprocessing_messages)
    assert mock_register_dcg_and_atlas.call_count == len(preprocessing_messages)
    assert mock_register_grid_square.call_count == len(preprocessing_messages)
    if ("gray" not in colors) or ("gray" in colors and len(colors) == 1):
        assert mock_align_and_merge_call.call_count == len(preprocessing_messages)
    else:
        assert mock_align_and_merge_call.call_count == len(preprocessing_messages) * 3


@pytest.mark.parametrize(
    "test_params",
    (
        # Reverse list order? | Colors
        (False, ["gray"]),
        (True, ["gray"]),
        (False, ["red", "green", "blue"]),
        (True, ["cyan", "magenta", "yellow"]),
        (False, ["gray", "red", "green", "blue"]),
        (True, ["gray", "cyan", "magenta", "yellow"]),
    ),
)
def test_run_with_db(
    mocker: MockerFixture,
    rsync_basepath: Path,
    mock_ispyb_credentials,
    murfey_db_session: SQLModelSession,
    ispyb_db_session: SQLAlchemySession,
    test_params: tuple[bool, list[str]],
):
    # Unpack test params
    (shuffle_message, colors) = test_params

    # Create a session to insert for this test
    murfey_session: MurfeyDB.Session = get_or_create_db_entry(
        murfey_db_session,
        MurfeyDB.Session,
        lookup_kwargs={
            "id": ExampleVisit.murfey_session_id + 1,
            "name": visit_name,
            "visit": visit_name,
            "instrument_name": ExampleVisit.instrument_name,
        },
    )

    # 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,
    )

    # Mock the align and merge workflow call
    mock_align_and_merge_call = mocker.patch(
        "murfey.workflows.clem.register_preprocessing_results.run_align_and_merge"
    )

    # Patch the TransportManager object in the workflows called
    from murfey.server.ispyb import TransportManager

    mocker.patch(
        "murfey.workflows.clem.register_preprocessing_results._transport_object",
        new=TransportManager("PikaTransport"),
    )
    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"),
    )

    # Run the function
    preprocessing_messages = generate_preprocessing_messages(
        rsync_basepath=rsync_basepath,
        session_id=murfey_session.id,
        colors=colors,
    )
    if shuffle_message:
        preprocessing_messages.reverse()
    for message in preprocessing_messages:
        result = run(
            message=message,
            murfey_db=murfey_db_session,
        )
        assert result == {"success": True}

    # Each message should call the align-and-merge workflow thrice
    # if gray and colour channels are both present
    if ("gray" not in colors) or ("gray" in colors and len(colors) == 1):
        assert mock_align_and_merge_call.call_count == len(preprocessing_messages)
    else:
        assert mock_align_and_merge_call.call_count == len(preprocessing_messages) * 3

    # Murfey's DataCollectionGroup should have an entry
    murfey_dcg_search = murfey_db_session.exec(
        sm_select(MurfeyDB.DataCollectionGroup).where(
            MurfeyDB.DataCollectionGroup.session_id == murfey_session.id
        )
    ).all()
    assert len(murfey_dcg_search) == 1

    # GridSquare entries should be half the initial number of entries due to overwrites
    murfey_gs_search = murfey_db_session.exec(
        sm_select(MurfeyDB.GridSquare).where(
            MurfeyDB.GridSquare.session_id == murfey_session.id
        )
    ).all()
    assert len(murfey_gs_search) == (len(preprocessing_messages) - 1) // 2

    # 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

    # Determine the color flags and collection mode
    color_flags = {
        COLOR_FLAGS_MURFEY_TO_ISPYB[flag]: int(value)
        for flag, value in _get_color_flags(colors).items()
    }
    collection_mode = _determine_collection_mode(colors)

    # Atlas color flags and collection mode should be set correctly
    ispyb_atlas = ispyb_atlas_search[0]
    for flag, value in color_flags.items():
        assert getattr(ispyb_atlas, flag) == value
    assert ispyb_atlas.mode == collection_mode

    # ISPyB's GrridSquare should have half the number of intiail entries
    ispyb_gs_search = (
        ispyb_db_session.execute(
            sa_select(ISPyBDB.GridSquare).where(
                ISPyBDB.GridSquare.atlasId == ispyb_atlas.atlasId
            )
        )
        .scalars()
        .all()
    )
    assert len(ispyb_gs_search) == (len(preprocessing_messages) - 1) // 2
    for gs in ispyb_gs_search:
        # Check that all entries point to the denoised images ("_Lng_LVCC")
        assert gs.gridSquareImage is not None and "_Lng_LVCC" in gs.gridSquareImage

        # Check that the GridSquare color flags and collection mode are set correctly
        for flag, value in color_flags.items():
            assert getattr(gs, flag) == value
        assert gs.mode == collection_mode

    murfey_db_session.close()
    ispyb_db_session.close()