Further Improve robustness of model classifier tests (#4450)

Author: chrishalcrow

src/spikeinterface/curation/tests/common.py

@@ -76,8 +76,8 @@ def sorting_analyzer_for_unitrefine_curation():
     """Makes an analyzer whose first 10 units are good normal units, and 10 which are noise. We make them
     noise by using a spike trains which are uncorrelated with the recording for `sorting2`."""
 
-    recording, sorting_1 = generate_ground_truth_recording(num_channels=4, seed=1, num_units=10)
-    _, sorting_2 = generate_ground_truth_recording(num_channels=4, seed=2, num_units=10)
+    recording, sorting_1 = generate_ground_truth_recording(num_channels=4, seed=1, num_units=6)
+    _, sorting_2 = generate_ground_truth_recording(num_channels=4, seed=2, num_units=6)
     both_sortings = aggregate_units([sorting_1, sorting_2])
     analyzer = create_sorting_analyzer(sorting=both_sortings, recording=recording)
     analyzer.compute(["random_spikes", "noise_levels", "templates"])
@@ -113,9 +113,9 @@ def trained_pipeline_path(sorting_analyzer_for_unitrefine_curation):
             }
         )
         train_model(
-            analyzers=[analyzer],
+            analyzers=[analyzer] * 5,
             folder=trained_model_folder,
-            labels=[[1] * 10 + [0] * 10],
+            labels=[[1] * 6 + [0] * 6] * 5,
             imputation_strategies=["median"],
             scaling_techniques=["standard_scaler"],
             classifiers=["RandomForestClassifier"],

src/spikeinterface/curation/tests/test_model_based_curation.py

@@ -132,11 +132,11 @@ def test_model_based_classification_predict_labels(sorting_analyzer_for_unitrefi
     classified_units = model_based_classification.predict_labels()
     predictions = classified_units["prediction"].values
 
-    expected_result = np.array([1] * 10 + [0] * 10)
+    expected_result = np.array([1] * 6 + [0] * 6)
     assert np.all(predictions == expected_result)
 
     conversion = {0: "noise", 1: "good"}
-    expected_result_converted = np.array(["good"] * 10 + ["noise"] * 10)
+    expected_result_converted = np.array(["good"] * 6 + ["noise"] * 6)
     classified_units_labelled = model_based_classification.predict_labels(label_conversion=conversion)
     predictions_labelled = classified_units_labelled["prediction"]
     assert np.all(predictions_labelled == expected_result_converted)