1 files changed, 44 insertions, 43 deletions

diff --git a/test/sca/test_stacked_combine.py b/test/sca/test_stacked_combine.py
index 7941b46..5ef355c 100644
--- a/test/sca/test_stacked_combine.py
+++ b/test/sca/test_stacked_combine.py
@@ -2,18 +2,14 @@ import pytest
 from numba import cuda
 
 import numpy as np
-from pyecsca.sca import (
-    StackedTraces,
-    GPUTraceManager,
-    CombinedTrace
-)
+from pyecsca.sca import StackedTraces, GPUTraceManager, CombinedTrace, CPUTraceManager
 
 TPB = 128
-TRACE_COUNT = 2 ** 10
-TRACE_LEN = 2 ** 15
+TRACE_COUNT = 2**10
+TRACE_LEN = 2**10
 RTOL = 1e-5
 ATOL = 1e-5
-CHUNK_SIZE = 2 ** 5
+CHUNK_SIZE = 2**5
 STREAM_COUNT = 4
 
 
@@ -25,79 +21,84 @@ def samples():
 
 class Base:
     @pytest.fixture()
-    def gpu_manager(self):
+    def manager(self, samples):
         raise NotImplementedError("Subclasses should implement this")
 
-    def test_average(self, samples, gpu_manager):
-        avg_trace = gpu_manager.average()
+    def test_average(self, samples, manager):
+        avg_trace = manager.average()
         avg_cmp: np.ndarray = np.average(samples, 0)
 
         assert isinstance(avg_trace, CombinedTrace)
-        assert avg_trace.samples.shape == \
-            avg_cmp.shape
+        assert avg_trace.samples.shape == avg_cmp.shape
         assert all(np.isclose(avg_trace.samples, avg_cmp))
 
-    def test_standard_deviation(self, samples, gpu_manager):
-        std_trace = gpu_manager.standard_deviation()
+    def test_standard_deviation(self, samples, manager):
+        std_trace = manager.standard_deviation()
         std_cmp: np.ndarray = np.std(samples, 0)
 
         assert isinstance(std_trace, CombinedTrace)
-        assert std_trace.samples.shape == \
-            std_cmp.shape
+        assert std_trace.samples.shape == std_cmp.shape
         assert all(np.isclose(std_trace.samples, std_cmp))
 
-    def test_variance(self, samples, gpu_manager):
-        var_trace = gpu_manager.variance()
+    def test_variance(self, samples, manager):
+        var_trace = manager.variance()
         var_cmp: np.ndarray = np.var(samples, 0)
 
         assert isinstance(var_trace, CombinedTrace)
-        assert var_trace.samples.shape == \
-            var_cmp.shape
+        assert var_trace.samples.shape == var_cmp.shape
         assert all(np.isclose(var_trace.samples, var_cmp))
 
-    def test_average_and_variance(self, samples, gpu_manager):
-        avg_trace, var_trace = gpu_manager.average_and_variance()
+    def test_average_and_variance(self, samples, manager):
+        avg_trace, var_trace = manager.average_and_variance()
         avg_cmp: np.ndarray = np.average(samples, 0)
         var_cmp: np.ndarray = np.var(samples, 0)
 
         assert isinstance(avg_trace, CombinedTrace)
         assert isinstance(var_trace, CombinedTrace)
-        assert avg_trace.samples.shape == \
-            avg_cmp.shape
-        assert var_trace.samples.shape == \
-            var_cmp.shape
+        assert avg_trace.samples.shape == avg_cmp.shape
+        assert var_trace.samples.shape == var_cmp.shape
         assert all(np.isclose(avg_trace.samples, avg_cmp))
         assert all(np.isclose(var_trace.samples, var_cmp))
 
-    def test_pearson_coef(self, samples, gpu_manager):
+    def test_pearson_coef(self, samples, manager):
         np.random.seed(0x1234)
         intermediate_values = np.random.rand(TRACE_COUNT)
-        corr_gpu = gpu_manager.pearson_corr(intermediate_values)
-        corr_cmp = np.corrcoef(
-            samples, intermediate_values, rowvar=False)[-1, :-1]
+        corr_gpu = manager.pearson_corr(intermediate_values)
+        corr_cmp = np.corrcoef(samples, intermediate_values, rowvar=False)[-1, :-1]
 
         assert isinstance(corr_gpu, CombinedTrace)
-        assert corr_gpu.samples.shape == \
-            corr_cmp.shape
+        assert corr_gpu.samples.shape == corr_cmp.shape
 
-        assert all(np.isclose(corr_gpu.samples,
-                   corr_cmp, rtol=RTOL, atol=ATOL))
+        assert all(np.isclose(corr_gpu.samples, corr_cmp, rtol=RTOL, atol=ATOL))
 
+    def test_pearson_coef_invalid(self, samples, manager):
+        intermediate_values = np.ones(TRACE_COUNT)
+        with pytest.raises(ValueError):
+            manager.pearson_corr(intermediate_values)
 
-class TestNonChunked(Base):
+
+class TestGPUNonChunked(Base):
     @pytest.fixture()
-    def gpu_manager(self, samples):
+    def manager(self, samples):
         if not cuda.is_available():
             pytest.skip("CUDA not available")
         return GPUTraceManager(StackedTraces(samples), TPB)
 
 
-class TestChunked(Base):
+class TestGPUChunked(Base):
     @pytest.fixture()
-    def gpu_manager(self, samples):
+    def manager(self, samples):
         if not cuda.is_available():
             pytest.skip("CUDA not available")
-        return GPUTraceManager(StackedTraces(samples),
-                               TPB,
-                               chunk_size=CHUNK_SIZE,
-                               stream_count=STREAM_COUNT)
+        return GPUTraceManager(
+            StackedTraces(samples),
+            TPB,
+            chunk_size=CHUNK_SIZE,
+            stream_count=STREAM_COUNT,
+        )
+
+
+class TestCPU(Base):
+    @pytest.fixture()
+    def manager(self, samples):
+        return CPUTraceManager(StackedTraces(samples))