Tests for registering grid squares

Author: stephen-riggs

tests/workflows/spa/test_flush_spa_preprocess.py

@@ -0,0 +1,166 @@
+import os
+from unittest import mock
+
+import pytest
+from sqlalchemy import Session, create_engine, select
+
+from murfey.util.db import DataCollectionGroup, GridSquare, clear, setup
+from murfey.util.models import GridSquareParameters
+from murfey.workflows.spa import flush_spa_preprocess
+
+url = (
+    f"postgresql+psycopg2://{os.environ['POSTGRES_USER']}:{os.environ['POSTGRES_PASSWORD']}"
+    f"@{os.environ['POSTGRES_HOST']}:{os.environ['POSTGRES_PORT']}/{os.environ['POSTGRES_DB']}"
+)
+
+
+@pytest.fixture
+def start_postgres():
+    clear(url)
+    setup(url)
+
+
+engine = create_engine(url)
+
+
+@mock.patch("murfey.workflows.spa.flush_spa_preprocess._transport_object")
+def test_register_grid_square_update_add_locations(mock_transport, start_postgres):
+    """Test the updating of an existing grid square"""
+    # Create a grid square to update
+    grid_square = GridSquare(
+        id=1,
+        name=101,
+        session_id=2,
+        tag="session_tag",
+    )
+    with Session(engine) as murfey_db:
+        murfey_db.add(grid_square)
+        murfey_db.commit()
+
+    # Parameters to update with
+    new_parameters = GridSquareParameters(
+        tag="session_tag",
+        x_location=1.1,
+        y_location=1.2,
+        x_stage_position=1.3,
+        y_stage_position=1.4,
+    )
+
+    # Run the registration
+    with Session(engine) as murfey_db:
+        flush_spa_preprocess.register_grid_square(2, 101, new_parameters, murfey_db)
+
+    # Check this would have updated ispyb
+    mock_transport.do_update_grid_square.assert_called_with(101, new_parameters)
+
+    # Confirm the database was updated
+    with Session(engine) as murfey_db:
+        grid_square_final_parameters = murfey_db.exec(select(GridSquare)).one()
+    assert grid_square_final_parameters.x_location == new_parameters.x_location
+    assert grid_square_final_parameters.y_location == new_parameters.y_location
+    assert (
+        grid_square_final_parameters.x_stage_position == new_parameters.x_stage_position
+    )
+    assert (
+        grid_square_final_parameters.y_stage_position == new_parameters.y_stage_position
+    )
+
+
+@mock.patch("murfey.workflows.spa.flush_spa_preprocess._transport_object")
+def test_register_grid_square_update_add_nothing(mock_transport, start_postgres):
+    """Test the updating of an existing grid square, but with nothing to update with"""
+    # Create a grid square to update
+    grid_square = GridSquare(
+        id=1,
+        name=101,
+        session_id=2,
+        tag="session_tag",
+        x_location=0.1,
+        y_location=0.2,
+        x_stage_location=0.3,
+        y_stage_location=0.4,
+    )
+    with Session(engine) as murfey_db:
+        murfey_db.add(grid_square)
+        murfey_db.commit()
+
+    # Parameters to update with
+    new_parameters = GridSquareParameters(tag="session_tag")
+
+    # Run the registration
+    with Session(engine) as murfey_db:
+        flush_spa_preprocess.register_grid_square(2, 101, new_parameters, murfey_db)
+
+    # Check this would have updated ispyb
+    mock_transport.do_update_grid_square.assert_called_with(101, new_parameters)
+
+    # Confirm the database was not updated
+    with Session(engine) as murfey_db:
+        grid_square_final_parameters = murfey_db.exec(select(GridSquare)).one()
+    assert grid_square_final_parameters.x_location == grid_square.x_location
+    assert grid_square_final_parameters.y_location == grid_square.y_location
+    assert grid_square_final_parameters.x_stage_position == grid_square.x_stage_position
+    assert grid_square_final_parameters.y_stage_position == grid_square.y_stage_position
+
+
+@mock.patch("murfey.workflows.spa.flush_spa_preprocess._transport_object")
+def test_register_grid_square_insert_with_ispyb(
+    mock_transport, start_postgres, tmp_path
+):
+    # Create a data collection group for lookups
+    grid_square = DataCollectionGroup(
+        id=1,
+        session_id=2,
+        tag="session_tag",
+        atlas_id=90,
+    )
+    with Session(engine) as murfey_db:
+        murfey_db.add(grid_square)
+        murfey_db.commit()
+
+    # Set the ispyb return
+    mock_transport.do_insert_grid_square.return_value = {
+        "return_value": 1,
+        "success": True,
+    }
+
+    # Parameters to update with
+    new_parameters = GridSquareParameters(
+        tag="session_tag",
+        x_location=1.1,
+        y_location=1.2,
+        x_stage_position=1.3,
+        y_stage_position=1.4,
+        readout_area_x=2048,
+        readout_area_y=1024,
+        thumbnail_size_x=512,
+        thumbnail_size_y=256,
+        pixel_size=1.02,
+        image=f"{tmp_path}/image_path",
+        angle=12.5,
+    )
+
+    # Run the registration
+    with Session(engine) as murfey_db:
+        flush_spa_preprocess.register_grid_square(2, 101, new_parameters, murfey_db)
+
+    # Check this would have updated ispyb
+    mock_transport.do_insert_grid_square.assert_called_with(90, 101, new_parameters)
+
+    # Confirm the database entry was made
+    with Session(engine) as murfey_db:
+        grid_square_final_parameters = murfey_db.exec(select(GridSquare)).one()
+    assert grid_square_final_parameters.id == 1
+    assert grid_square_final_parameters.name == 101
+    assert grid_square_final_parameters.session_id == 2
+    assert grid_square_final_parameters.tag == "session_tag"
+    assert grid_square_final_parameters.x_location == 1.1
+    assert grid_square_final_parameters.y_location == 1.2
+    assert grid_square_final_parameters.x_stage_position == 1.3
+    assert grid_square_final_parameters.y_stage_position == 1.4
+    assert grid_square_final_parameters.readout_area_x == 2048
+    assert grid_square_final_parameters.readout_area_y == 1024
+    assert grid_square_final_parameters.thumbnail_size_x == 512
+    assert grid_square_final_parameters.thumbnail_size_y == 256
+    assert grid_square_final_parameters.pixel_size == 1.02
+    assert grid_square_final_parameters.image == f"{tmp_path}/image_path"