Update tests for remove_bad_channels, check for number of segments  and their sizes being maintained after the preprocessor.
@@ -16,7 +16,7 @@ def test_remove_bad_channels():
sampling_frequency = 30000.
durations = [10.325, 3.5]
- num_segments = len(durations)
+ num_segments = len(durations)j
num_timepoints = [int(sampling_frequency * d) for d in durations]
traces_list = []
@@ -35,7 +35,12 @@ def test_remove_bad_channels():
rec2 = remove_bad_channels(rec, bad_threshold=5.)
+ # Check that the noisy channel is taken out
assert np.array_equal(rec2.get_channel_ids(), [0,2,3])
+ # Check that the number of segments is maintained after preprocessor
+ assert np.array_equal(rec2.get_num_segments(), rec.get_num_segments())
+ # Check that the size of the segments os maintained after preprocessor
+ assert np.array_equal(*([r.get_num_frames(x) for x in range(rec.get_num_segments())] for r in [rec, rec2]))
if __name__ == '__main__':
