Made tests use temporary directory. Added missing unit tests.

Author: MichaelClerx

tests/test_trace_class.py

@@ -1,32 +1,33 @@
+#!/usr/bin/env python
 import json
 import os
+import tempfile
 import unittest
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 
-from syncropatch_export.trace import Trace as tr
+from syncropatch_export.trace import Trace
 from syncropatch_export.voltage_protocols import VoltageProtocol
 
 
 class TestTraceClass(unittest.TestCase):
-    def setUp(self):
-        filepath = os.path.join('tests', 'test_data', '13112023_MW2_FF',
-                                'staircaseramp (2)_2kHz_15.01.07')
-        json_file = "staircaseramp (2)_2kHz_15.01.07"
+    """
+    Tests both the Trace and VoltageProtocol classes.
+    """
 
-        self.output_dir = os.path.join('test_output', 'test_trace_class')
 
-        if not os.path.exists(self.output_dir):
-            os.makedirs(self.output_dir)    # pragma: no-cover
-        self.test_trace = tr(filepath, json_file)
+    def setUp(self):
+        f = 'staircaseramp (2)_2kHz_15.01.07'
+        self.trace = Trace(
+            os.path.join('tests', 'test_data', '13112023_MW2_FF', f), f)
 
     def test_protocol_descriptions(self):
-        voltages = self.test_trace.get_voltage()
-        times = self.test_trace.get_times()
+        voltages = self.trace.get_voltage()
+        times = self.trace.get_times()
 
-        protocol_from_json = self.test_trace.get_voltage_protocol()
+        protocol_from_json = self.trace.get_voltage_protocol()
         holding_potential = protocol_from_json.get_holding_potential()
         protocol_desc = VoltageProtocol.from_voltage_trace(voltages, times,
                                                            holding_potential)
@@ -43,18 +44,17 @@ def test_protocol_descriptions(self):
         self.assertLess(v_error, 1e-4)
 
     def test_protocol_export(self):
-        protocol = self.test_trace.get_voltage_protocol()
-        protocol.export_txt(os.path.join(self.output_dir, 'protocol.txt'))
-        json_protocol = self.test_trace.get_voltage_protocol_json()
-
-        with open(os.path.join(self.output_dir, 'protocol.json'), 'w') as fin:
-            json.dump(json_protocol, fin)
+        with tempfile.TemporaryDirectory() as d:
+            protocol = self.trace.get_voltage_protocol()
+            protocol.export_txt(os.path.join(d, 'protocol.txt'))
+            json_protocol = self.trace.get_voltage_protocol_json()
+            with open(os.path.join(d, 'protocol.json'), 'w') as fin:
+                json.dump(json_protocol, fin)
 
     def test_protocol_timeseries(self):
-        voltages = self.test_trace.get_voltage()
-        times = self.test_trace.get_times()
-
-        voltage_protocol = self.test_trace.get_voltage_protocol()
+        voltages = self.trace.get_voltage()
+        times = self.trace.get_times()
+        voltage_protocol = self.trace.get_voltage_protocol()
 
         def voltage_func(t):
             for tstart, tend, vstart, vend in voltage_protocol.get_all_sections():
@@ -68,33 +68,46 @@ def voltage_func(t):
         for t, v in zip(times, voltages):
             self.assertLess(voltage_func(t) - v, 1e-3)
 
+    def test_protocol_get_step_start_times(self):
+        a = list(self.trace.get_voltage_protocol().get_step_start_times())
+        b = [0, 250, 300, 696, 896, 1896, 2396, 3396, 3896, 4396, 4896, 5396,
+             5896, 6396, 6896, 7396, 7896, 8396, 8896, 9396, 9896, 10396,
+             10896, 11396, 11896, 12396, 12896, 13896, 14396, 14406, 14502,
+             14892]
+        self.assertEqual(a, b)
+
+    def test_protocol_get_ramps(self):
+        a = np.array(self.trace.get_voltage_protocol().get_ramps())
+        b = np.array([[300, 696, -120, -80], [14406, 14502, -70, -110]])
+        self.assertEqual(a.shape, b.shape)
+        self.assertTrue(np.all(a == b))
+
     def test_get_QC(self):
-        tr = self.test_trace
-        QC_values = tr.get_onboard_QC_values()
+        QC_values = self.trace.get_onboard_QC_values()
         self.assertGreater(len(QC_values), 0)
-        df = tr.get_onboard_QC_df()
+        df = self.trace.get_onboard_QC_df()
 
         self.assertGreater(df.shape[0], 0)
         self.assertGreater(df.shape[1], 0)
 
     def test_get_traces(self):
-        tr = self.test_trace
-        v = tr.get_voltage()
-        ts = tr.get_times()
-        all_traces = tr.get_all_traces(leakcorrect=True)
-        all_traces = tr.get_all_traces()
+        v = self.trace.get_voltage()
+        ts = self.trace.get_times()
+        all_traces = self.trace.get_all_traces(leakcorrect=True)
+        all_traces = self.trace.get_all_traces()
 
         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)))
 
+        '''
         if self.output_dir:
             # plot test output
             fig, (ax1, ax2) = plt.subplots(2, 1)
             ax1.set_title("Example Sweeps")
-            some_sweeps = tr.get_trace_sweeps([0])['A01']
+            some_sweeps = self.trace.get_trace_sweeps([0])['A01']
 
             ax1.plot(ts, np.transpose(some_sweeps), color='grey', alpha=0.5)
             ax1.set_ylabel('Current')
@@ -104,13 +117,13 @@ def test_get_traces(self):
             ax2.set_ylabel('Voltage')
             ax2.set_xlabel('Time')
             plt.tight_layout()
-            plt.savefig(os.path.join(self.output_dir,
-                                     'example_trace'))
+            plt.savefig(os.path.join(self.output_dir, 'example_trace'))
             plt.close(fig)
+        '''
 
     def test_qc_df(self):
-        dfs = [self.test_trace.get_onboard_QC_df(sweeps=[0]),
-               self.test_trace.get_onboard_QC_df(sweeps=None)]
+        dfs = [self.trace.get_onboard_QC_df(sweeps=[0]),
+               self.trace.get_onboard_QC_df(sweeps=None)]
         for res in dfs:
             # Check res is a pd.DataFrame
             self.assertIsInstance(res, pd.DataFrame)
@@ -125,3 +138,6 @@ def test_qc_df(self):
         # Check restricting number of sweeps returns less data
         self.assertLess(dfs[0].shape[0], dfs[1].shape[0])
 
+
+if __name__ == '__main__':
+    unittest.main()