#21 Update test_qc2

pcpostprocess/hergQC.py

@@ -461,7 +461,8 @@ def qc6(self, recording1, win=None, label=''):
         val = np.mean(recording1[i:f])
 
         if self.plot_dir and self._debug:
-            plt.axvspan(win[0], win[1], color='grey', alpha=.1)
+            if win is not None:
+                plt.axvspan(win[0], win[1], color='grey', alpha=.1)
             plt.plot(recording1, label='recording1')
             plt.savefig(os.path.join(self.plot_dir, f"qc6_{label}"))
             plt.clf()

tests/test_herg_qc.py

@@ -104,21 +104,17 @@ def test_qc2(self):
                         voltage=self.voltage)
 
         # qc2 checks that raw and subtracted SNR are above a minimum threshold
-        recording = np.asarray([0, 1] * 500 + [0, 10] * 500)  # snr = 70.75
-        result = hergqc.qc2(recording)
-        self.assertTrue(result[0], f"({result[1]})")
-
-        recording = np.asarray([0, 1] * 500 + [0, 6.03125] * 500)  # snr = 25.02
-        result = hergqc.qc2(recording)
-        self.assertTrue(result[0], f"({result[1]})")
-
-        recording = np.asarray([0, 1] * 500 + [0, 6.015625] * 500)  # snr = 24.88
-        result = hergqc.qc2(recording)
-        self.assertFalse(result[0], f"({result[1]})")
+        test_matrix = [
+            (10, True),  # snr = 70.75
+            (6.03125, True),  # snr = 25.02
+            (6.015625, False),  # snr = 24.88
+            (1, False),  # snr = 1.0
+        ]
 
-        recording = np.asarray([0, 1] * 1000)  # snr = 1.0
-        result = hergqc.qc2(recording)
-        self.assertFalse(result[0], f"({result[1]})")
+        for i, expected in test_matrix:
+            recording = np.asarray([0, 1] * 500 + [0, i] * 500)  # snr = 70.75
+            result = hergqc.qc2(recording)
+            self.assertEqual(result[0], expected, f"({i}: {result[1]})")
 
         # TODO: Test on select data
 
@@ -200,6 +196,7 @@ def test_qc5(self):
             (0.5, False),
             (0.75, False),
             (1.0, False),
+            (2.0, False),
         ]
 
         for i, expected in test_matrix: