Add tests for trim_waveform_data

tests/test_model/test_handling_data.py

@@ -13,12 +13,103 @@ from sohstationviewer.model.handling_data import (
     trim_downsample_WFChan,
     get_eachDay5MinList,
     get_trimTPSData,
+    trim_waveform_data,
 )
 
 ORIGINAL_CHAN_SIZE_LIMIT = const.CHAN_SIZE_LIMIT
 ORIGINAL_RECAL_SIZE_LIMIT = const.RECAL_SIZE_LIMIT
 
 
+class TestTrimWfData(TestCase):
+    def setUp(self) -> None:
+        self.channel_data = {}
+        self.traces_info = []
+        self.channel_data['tracesInfo'] = self.traces_info
+
+        for i in range(100):
+            trace_size = 100
+            start_time = i * trace_size
+            trace = {}
+            trace['startTmEpoch'] = start_time
+            trace['endTmEpoch'] = start_time + trace_size - 1
+            self.traces_info.append(trace)
+        self.start_time = 2500
+        self.end_time = 7500
+
+    def test_data_is_trimmed_neither_start_nor_end_time_is_trace_start_or_end_time(self):  # noqa: E501
+        self.start_time = 2444
+        self.end_time = 7444
+        trimmed_traces_list = trim_waveform_data(
+            self.channel_data, self.start_time, self.end_time
+        )
+        
+        self.assertTrue(
+            trimmed_traces_list[0]['startTmEpoch'] <= self.start_time)
+        self.assertTrue(
+            trimmed_traces_list[0]['endTmEpoch'] > self.start_time
+        )
+        trimmed_traces_list.pop(0)
+        trimmed_traces_list.pop()
+        is_left_trimmed = all(trace['startTmEpoch'] > self.start_time
+                              for trace in trimmed_traces_list)
+        is_right_trimmed = all(trace['endTmEpoch'] <= self.end_time
+                               for trace in trimmed_traces_list)
+        self.assertTrue(is_left_trimmed and is_right_trimmed)
+
+    def test_data_out_of_range(self):
+        with self.subTest('test_start_time_later_than_data_end_time'):
+            self.start_time = 12500
+            self.end_time = 17500
+            self.assertFalse(
+                trim_downsample_WFChan(self.channel_data, self.start_time,
+                                       self.end_time, True)
+            )
+        with self.subTest('test_end_time_earlier_than_data_start_time'):
+            self.start_time = -7500
+            self.end_time = -2500
+            self.assertFalse(
+                trim_downsample_WFChan(self.channel_data, self.start_time,
+                                       self.end_time, True)
+            )
+
+    def test_no_data(self):
+        self.channel_data['tracesInfo'] = []
+        with self.assertRaises(IndexError):
+            trim_waveform_data(
+                self.channel_data, self.start_time, self.end_time
+            )
+
+    def test_end_time_earlier_than_start_time(self):
+        self.start_time, self.end_time = self.end_time, self.start_time
+        trimmed_traces_list = trim_waveform_data(
+            self.channel_data, self.start_time, self.end_time
+        )
+        self.assertListEqual(trimmed_traces_list, [])
+
+    def test_data_does_not_need_to_be_trimmed(self):
+        with self.subTest('test_start_time_earlier_than_trace_earliest_time'):
+            self.start_time = -2500
+            self.end_time = 7500
+            trimmed_traces_list = trim_waveform_data(
+                self.channel_data, self.start_time, self.end_time
+            )
+            self.assertEqual(len(trimmed_traces_list), 76)
+        with self.subTest('test_end_time_later_than_trace_latest_time'):
+            self.start_time = 2500
+            self.end_time = 12500
+            trimmed_traces_list = trim_waveform_data(
+                self.channel_data, self.start_time, self.end_time
+            )
+            self.assertEqual(len(trimmed_traces_list), 75)
+        with self.subTest('test_data_contained_in_time_range'):
+            self.start_time = self.traces_info[0]['startTmEpoch']
+            self.end_time = self.traces_info[-1]['endTmEpoch']
+            trimmed_traces_list = trim_waveform_data(
+                self.channel_data, self.start_time, self.end_time
+            )
+            self.assertEqual(len(trimmed_traces_list), len(self.traces_info))
+
+
 class TestTrimDownsampleWfChan(TestCase):
     def no_file_memmap(self, file_path: Path, **kwargs):
         # Data will look the same as times. This has two benefits:
@@ -60,22 +151,6 @@ class TestTrimDownsampleWfChan(TestCase):
         self.start_time = 2500
         self.end_time = 7500
 
-    def test_data_out_of_range(self):
-        with self.subTest('test_start_time_later_than_data_end_time'):
-            self.start_time = 12500
-            self.end_time = 17500
-            self.assertFalse(
-                trim_downsample_WFChan(self.channel_data, self.start_time,
-                                       self.end_time, True)
-            )
-        with self.subTest('test_end_time_earlier_than_data_start_time'):
-            self.start_time = -7500
-            self.end_time = -2500
-            self.assertFalse(
-                trim_downsample_WFChan(self.channel_data, self.start_time,
-                                       self.end_time, True)
-            )
-
     def test_result_is_stored(self):
         trim_downsample_WFChan(self.channel_data, self.start_time,
                                self.end_time, True)
@@ -84,16 +159,6 @@ class TestTrimDownsampleWfChan(TestCase):
         self.assertTrue('data' in self.channel_data)
         self.assertGreater(len(self.channel_data['data']), 0)
 
-    def test_data_is_trimmed(self):
-        trim_downsample_WFChan(self.channel_data, self.start_time,
-                               self.end_time, True)
-        is_times_trimmed = ((self.channel_data['times'] >= self.start_time) &
-                            (self.channel_data['times'] <= self.end_time))
-        self.assertTrue(is_times_trimmed.all())
-        is_data_trimmed = ((self.channel_data['data'] >= self.start_time) &
-                           (self.channel_data['data'] <= self.end_time))
-        self.assertTrue(is_data_trimmed.all())
-
     @patch('sohstationviewer.model.handling_data.downsample', wraps=downsample)
     def test_data_is_downsampled(self, mock_downsample):
         const.CHAN_SIZE_LIMIT = 1000