Added missing unit tests

Author: MichaelClerx

syncropatch_export/trace.py

@@ -175,7 +175,7 @@ def get_trace_sweeps(self, sweeps=None, leakcorrect=False):
         corrected data can be obtained by setting ``leakcorrect=True``.
 
         Args:
-            sweeps (int): The number of sweeps to return.
+            sweeps (list): A list of sweep indexes to return, e.g. ``[0, 1, 2]``.
             leakcorrect (bool): Used to choose corrected or uncorrected data.
 
         Returns:
@@ -191,18 +191,22 @@ def get_trace_sweeps(self, sweeps=None, leakcorrect=False):
             for ijWell in iCol:
                 out_dict[ijWell] = []
 
+        # No sweeps selected? Then return full set
         if sweeps is None:
-            # Sometimes NofSweeps seems to be incorrect
             sweeps = list(range(self.NofSweeps))
-
-        # Check `sweeps` is something sensible
-        elif len(sweeps) > self.NofSweeps:
-            raise ValueError('Required #sweeps > total #sweeps.')
-
-        # convert negative values to positive
-        for i, sweep in enumerate(sweeps):
-            if sweep < 0:
-                sweeps[i] = self.NofSweeps + sweep
+        else:
+            # Allow negative values to index later sweeps
+            sweeps = [self.NofSweeps + x if x < 0 else x for x in sweeps]
+            # Check all sweeps exist
+            if max(sweeps) >= self.NofSweeps:
+                raise ValueError(
+                    f'Invalid sweep selection: sweep {max(sweeps)} requested,'
+                    f' but only {self.NofSweeps} available.')
+            if min(sweeps) < 0:
+                raise ValueError(
+                    f'Invalid sweep selection: sweep'
+                    f' {min(sweeps) - self.NofSweeps} requested, but only'
+                    f' {self.NofSweeps} available.')
 
         trace_file_idxs, idx_is = self.get_trace_file(sweeps)
 

tests/test_trace_class.py

@@ -99,13 +99,26 @@ def test_get_traces(self):
         ts = self.trace.get_times()
         all_traces = self.trace.get_all_traces(leakcorrect=True)
         all_traces = self.trace.get_all_traces()
+        # TODO: Check the output, numerically, by comparing a few points
 
         self.assertTrue(np.all(np.isfinite(v)))
         self.assertTrue(np.all(np.isfinite(ts)))
 
         for well, trace in all_traces.items():
             self.assertTrue(np.all(np.isfinite(trace)))
 
+        # Test complex sweep selection
+        a = self.trace.get_trace_sweeps([-1, -2])
+        b = self.trace.get_trace_sweeps([1, 0])
+        self.assertEqual(len(a), len(b))
+        self.assertTrue(np.all(a['A01'] == b['A01']))
+
+        # Test asking for non-existent sweeps
+        self.assertRaisesRegex(ValueError, 'Invalid sweep selection',
+                               self.trace.get_trace_sweeps, [2])
+        self.assertRaisesRegex(ValueError, 'Invalid sweep selection',
+                               self.trace.get_trace_sweeps, [-3])
+
         '''
         # plot test output
         if False: