Redo the TraceSet IO interfaces.

1 files changed, 72 insertions, 21 deletions

diff --git a/test/sca/test_traceset.py b/test/sca/test_traceset.py
index 6d48707..a3ff526 100644
--- a/test/sca/test_traceset.py
+++ b/test/sca/test_traceset.py
@@ -1,8 +1,17 @@
 import os.path
+import shutil
 import tempfile
 from unittest import TestCase
 
-from pyecsca.sca import TraceSet, InspectorTraceSet, ChipWhispererTraceSet
+import numpy as np
+
+from pyecsca.sca import (TraceSet, InspectorTraceSet, ChipWhispererTraceSet, PickleTraceSet,
+                         HDF5TraceSet, Trace)
+
+EXAMPLE_TRACES = [Trace(np.array([20, 40, 50, 50, 10], dtype=np.dtype("i1")), None, None),
+                  Trace(np.array([1, 2, 3, 4, 5], dtype=np.dtype("i1")), None, None),
+                  Trace(np.array([6, 7, 8, 9, 10], dtype=np.dtype("i1")), None, None)]
+EXAMPLE_KWARGS = {"num_traces": 3, "thingy": "abc"}
 
 
 class TraceSetTests(TestCase):
@@ -11,12 +20,14 @@ class TraceSetTests(TestCase):
         self.assertIsNotNone(TraceSet())
         self.assertIsNotNone(InspectorTraceSet())
         self.assertIsNotNone(ChipWhispererTraceSet())
+        self.assertIsNotNone(PickleTraceSet())
+        self.assertIsNotNone(HDF5TraceSet())
 
 
 class InspectorTraceSetTests(TestCase):
 
     def test_load_fname(self):
-        result = InspectorTraceSet("test/data/example.trs")
+        result = InspectorTraceSet.read("test/data/example.trs")
         self.assertIsNotNone(result)
         self.assertEqual(result.global_title, "Example trace set")
         self.assertEqual(len(result), 10)
@@ -27,36 +38,76 @@ class InspectorTraceSetTests(TestCase):
 
     def test_load_file(self):
         with open("test/data/example.trs", "rb") as f:
-            self.assertIsNotNone(InspectorTraceSet(f))
+            self.assertIsNotNone(InspectorTraceSet.read(f))
 
     def test_load_bytes(self):
         with open("test/data/example.trs", "rb") as f:
-            self.assertIsNotNone(InspectorTraceSet(f.read()))
-
-    def test_get_bytes(self):
-        self.assertIsNotNone(bytes(InspectorTraceSet("test/data/example.trs")))
-
-    def test_keep_traces(self):
-        trace_set = InspectorTraceSet("test/data/example.trs")
-        self.assertIsNotNone(trace_set.raw)
-        trace_set = InspectorTraceSet("test/data/example.trs", keep_raw_traces=False)
-        self.assertIsNone(trace_set.raw)
+            self.assertIsNotNone(InspectorTraceSet.read(f.read()))
 
     def test_save(self):
-        trace_set = InspectorTraceSet("test/data/example.trs")
+        trace_set = InspectorTraceSet.read("test/data/example.trs")
         with tempfile.TemporaryDirectory() as dirname:
             path = os.path.join(dirname, "out.trs")
-            trace_set.save(path)
+            trace_set.write(path)
             self.assertTrue(os.path.exists(path))
-            self.assertIsNotNone(InspectorTraceSet(path))
+            self.assertIsNotNone(InspectorTraceSet.read(path))
 
-        with self.assertRaises(ValueError):
-            trace_set.save(None)
 
-
-class ChipWhispererTraceSetTest(TestCase):
+class ChipWhispererTraceSetTests(TestCase):
 
     def test_load_fname(self):
-        result = ChipWhispererTraceSet("test/data/", "chipwhisperer")
+        result = ChipWhispererTraceSet.read("test/data/config_chipwhisperer_.cfg")
         self.assertIsNotNone(result)
         self.assertEqual(len(result), 2)
+
+
+class PickleTraceSetTests(TestCase):
+
+    def test_load_fname(self):
+        result = PickleTraceSet.read("test/data/test.pickle")
+        self.assertIsNotNone(result)
+
+    def test_load_file(self):
+        with open("test/data/test.pickle", "rb") as f:
+            self.assertIsNotNone(PickleTraceSet.read(f))
+
+    def test_save(self):
+        trace_set = PickleTraceSet(*EXAMPLE_TRACES, **EXAMPLE_KWARGS)
+        with tempfile.TemporaryDirectory() as dirname:
+            path = os.path.join(dirname, "out.pickle")
+            trace_set.write(path)
+            self.assertTrue(os.path.exists(path))
+            self.assertIsNotNone(PickleTraceSet.read(path))
+
+
+class HDF5TraceSetTests(TestCase):
+
+    def test_load_fname(self):
+        result = HDF5TraceSet.read("test/data/test.h5")
+        self.assertIsNotNone(result)
+
+    def test_load_file(self):
+        with open("test/data/test.h5", "rb") as f:
+            self.assertIsNotNone(HDF5TraceSet.read(f))
+
+    def test_inplace(self):
+        with tempfile.TemporaryDirectory() as dirname:
+            path = os.path.join(dirname, "test.h5")
+            shutil.copy("test/data/test.h5", path)
+            trace_set = HDF5TraceSet.inplace(path)
+            self.assertIsNotNone(trace_set)
+            test_trace = Trace(np.array([6, 7], dtype=np.dtype("i1")), None, None, meta={"thing": "ring"})
+            trace_set[0] = test_trace
+            trace_set.save()
+            trace_set.close()
+
+            test_set = HDF5TraceSet.read(path)
+            self.assertEquals(test_set[0], test_trace)
+
+    def test_save(self):
+        trace_set = HDF5TraceSet(*EXAMPLE_TRACES, **EXAMPLE_KWARGS)
+        with tempfile.TemporaryDirectory() as dirname:
+            path = os.path.join(dirname, "out.h5")
+            trace_set.write(path)
+            self.assertTrue(os.path.exists(path))
+            self.assertIsNotNone(HDF5TraceSet.read(path))