Mocking count didn't work

Author: stephen-riggs

tests/workflows/tomo/test_tomo_picking.py

@@ -152,13 +152,11 @@ def test_picked_tomogram_not_run_class2d(
 
 @mock.patch("murfey.workflows.tomo.picking._transport_object")
 @mock.patch("murfey.workflows.tomo.picking._ids_tomo_classification")
-@mock.patch("murfey.workflows.tomo.picking.func")
 def test_picked_tomogram_run_class2d_with_diameter(
-    mock_func, mock_ids, mock_transport, murfey_db_session: Session, tmp_path
+    mock_ids, mock_transport, murfey_db_session: Session, tmp_path
 ):
     """Run the picker feedback with a pre-determined particle diameter"""
     mock_transport.feedback_queue = "murfey_feedback"
-    mock_func.count.return_value = 10001
 
     # Insert table dependencies
     dcg_id, dc_id, pj_id = set_up_picking_db(murfey_db_session)
@@ -175,7 +173,7 @@ def test_picked_tomogram_run_class2d_with_diameter(
             "particle_diameter": 200,
         },
     )
-    for particle in range(10):
+    for particle in range(10001):
         get_or_create_db_entry(
             murfey_db_session,
             ParticleSizes,
@@ -201,7 +199,6 @@ def test_picked_tomogram_run_class2d_with_diameter(
     picking._register_picked_tomogram_use_diameter(message, murfey_db_session)
 
     mock_ids.assert_called_once_with(0, "em-tomo-class2d", murfey_db_session)
-    mock_func.count.assert_called_once()
 
     tomograms_db = murfey_db_session.exec(
         select(TomogramPicks).where(TomogramPicks.pj_id == 1)
@@ -239,13 +236,11 @@ def test_picked_tomogram_run_class2d_with_diameter(
 
 @mock.patch("murfey.workflows.tomo.picking._transport_object")
 @mock.patch("murfey.workflows.tomo.picking._ids_tomo_classification")
-@mock.patch("murfey.workflows.tomo.picking.func")
 def test_picked_tomogram_run_class2d_estimate_diameter(
-    mock_func, mock_ids, mock_transport, murfey_db_session: Session, tmp_path
+    mock_ids, mock_transport, murfey_db_session: Session, tmp_path
 ):
     """Run the picker feedback for Class2D, including diameter estimation"""
     mock_transport.feedback_queue = "murfey_feedback"
-    mock_func.count.return_value = 10001
 
     # Insert table dependencies
     dcg_id, dc_id, pj_id = set_up_picking_db(murfey_db_session)
@@ -262,7 +257,7 @@ def test_picked_tomogram_run_class2d_estimate_diameter(
             "particle_diameter": None,
         },
     )
-    for particle in range(10):
+    for particle in range(10001):
         get_or_create_db_entry(
             murfey_db_session,
             ParticleSizes,
@@ -300,7 +295,6 @@ def test_picked_tomogram_run_class2d_estimate_diameter(
     picking._register_picked_tomogram_use_diameter(message, murfey_db_session)
 
     mock_ids.assert_called_once_with(0, "em-tomo-class2d", murfey_db_session)
-    mock_func.count.assert_called_once()
 
     # Two mock calls - one flushed tomogram and one new
     assert mock_transport.send.call_count == 2