1 files changed, 24 insertions, 0 deletions

diff --git a/test/utils.py b/test/utils.py
new file mode 100644
index 0000000..e7d5819
--- /dev/null
+++ b/test/utils.py
@@ -0,0 +1,24 @@
+import matplotlib.pyplot as plt
+from unittest import TestCase
+from pyecsca import Trace
+from os.path import join, exists
+from os import mkdir, getenv
+
+
+def slow(func):
+    func.slow = 1
+    return func
+
+
+def plot(case: TestCase, *traces: Trace):
+    if getenv("PYECSCA_TEST_PLOTS") is None:
+        return
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    for i, trace in enumerate(traces):
+        ax.plot(trace.samples, label=str(i))
+    ax.legend(loc="best")
+    directory = join("test", "plots")
+    if not exists(directory):
+        mkdir(directory)
+    plt.savefig(join(directory, case.id() + ".png"))