Replaced use of 'start_postgres' with test-safe 'murfey_db_session' fixtures

Author: tieneupin

tests/workflows/spa/test_flush_spa_preprocess.py

@@ -5,11 +5,12 @@
 from murfey.util.db import DataCollectionGroup, GridSquare
 from murfey.util.models import GridSquareParameters
 from murfey.workflows.spa import flush_spa_preprocess
-from tests import murfey_db_engine
 
 
 @mock.patch("murfey.workflows.spa.flush_spa_preprocess._transport_object")
-def test_register_grid_square_update_add_locations(mock_transport, start_postgres):
+def test_register_grid_square_update_add_locations(
+    mock_transport, murfey_db_session: Session
+):
     """Test the updating of an existing grid square"""
     # Create a grid square to update
     grid_square = GridSquare(
@@ -18,9 +19,8 @@ def test_register_grid_square_update_add_locations(mock_transport, start_postgre
         session_id=2,
         tag="session_tag",
     )
-    with Session(murfey_db_engine) as murfey_db:
-        murfey_db.add(grid_square)
-        murfey_db.commit()
+    murfey_db_session.add(grid_square)
+    murfey_db_session.commit()
 
     # Parameters to update with
     new_parameters = GridSquareParameters(
@@ -32,15 +32,13 @@ def test_register_grid_square_update_add_locations(mock_transport, start_postgre
     )
 
     # Run the registration
-    with Session(murfey_db_engine) as murfey_db:
-        flush_spa_preprocess.register_grid_square(2, 101, new_parameters, murfey_db)
+    flush_spa_preprocess.register_grid_square(2, 101, new_parameters, murfey_db_session)
 
     # Check this would have updated ispyb
     mock_transport.do_update_grid_square.assert_called_with(1, new_parameters)
 
     # Confirm the database was updated
-    with Session(murfey_db_engine) as murfey_db:
-        grid_square_final_parameters = murfey_db.exec(select(GridSquare)).one()
+    grid_square_final_parameters = murfey_db_session.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 (
@@ -52,7 +50,9 @@ def test_register_grid_square_update_add_locations(mock_transport, start_postgre
 
 
 @mock.patch("murfey.workflows.spa.flush_spa_preprocess._transport_object")
-def test_register_grid_square_update_add_nothing(mock_transport, start_postgres):
+def test_register_grid_square_update_add_nothing(
+    mock_transport, murfey_db_session: Session
+):
     """Test the updating of an existing grid square, but with nothing to update with"""
     # Create a grid square to update
     grid_square = GridSquare(
@@ -65,23 +65,20 @@ def test_register_grid_square_update_add_nothing(mock_transport, start_postgres)
         x_stage_position=0.3,
         y_stage_position=0.4,
     )
-    with Session(murfey_db_engine) as murfey_db:
-        murfey_db.add(grid_square)
-        murfey_db.commit()
+    murfey_db_session.add(grid_square)
+    murfey_db_session.commit()
 
     # Parameters to update with
     new_parameters = GridSquareParameters(tag="session_tag")
 
     # Run the registration
-    with Session(murfey_db_engine) as murfey_db:
-        flush_spa_preprocess.register_grid_square(2, 101, new_parameters, murfey_db)
+    flush_spa_preprocess.register_grid_square(2, 101, new_parameters, murfey_db_session)
 
     # Check this would have updated ispyb
     mock_transport.do_update_grid_square.assert_called_with(1, new_parameters)
 
     # Confirm the database was not updated
-    with Session(murfey_db_engine) as murfey_db:
-        grid_square_final_parameters = murfey_db.exec(select(GridSquare)).one()
+    grid_square_final_parameters = murfey_db_session.exec(select(GridSquare)).one()
     assert grid_square_final_parameters.x_location == 0.1
     assert grid_square_final_parameters.y_location == 0.2
     assert grid_square_final_parameters.x_stage_position == 0.3
@@ -90,7 +87,7 @@ def test_register_grid_square_update_add_nothing(mock_transport, start_postgres)
 
 @mock.patch("murfey.workflows.spa.flush_spa_preprocess._transport_object")
 def test_register_grid_square_insert_with_ispyb(
-    mock_transport, start_postgres, tmp_path
+    mock_transport, murfey_db_session, tmp_path
 ):
     # Create a data collection group for lookups
     grid_square = DataCollectionGroup(
@@ -99,9 +96,8 @@ def test_register_grid_square_insert_with_ispyb(
         tag="session_tag",
         atlas_id=90,
     )
-    with Session(murfey_db_engine) as murfey_db:
-        murfey_db.add(grid_square)
-        murfey_db.commit()
+    murfey_db_session.add(grid_square)
+    murfey_db_session.commit()
 
     # Set the ispyb return
     mock_transport.do_insert_grid_square.return_value = {
@@ -126,15 +122,13 @@ def test_register_grid_square_insert_with_ispyb(
     )
 
     # Run the registration
-    with Session(murfey_db_engine) as murfey_db:
-        flush_spa_preprocess.register_grid_square(2, 101, new_parameters, murfey_db)
+    flush_spa_preprocess.register_grid_square(2, 101, new_parameters, murfey_db_session)
 
     # 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(murfey_db_engine) as murfey_db:
-        grid_square_final_parameters = murfey_db.exec(select(GridSquare)).one()
+    grid_square_final_parameters = murfey_db_session.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