add more segement_axis tests
Tests for dit.inference.segmentaxis.
"""
+import pytest
+
import numpy as np
from dit.inference.segmentaxis import segment_axis
@@ -29,7 +31,7 @@ def test_sa2():
Test padding.
"""
a = np.arange(5)
- sa = segment_axis(a, 2, 0, end='pad', endvalue=7)
+ sa = segment_axis(a, 2, 0, axis=0, end='pad', endvalue=7)
sa_correct = np.array([[0, 1],
[2, 3],
[4, 7]])
@@ -57,3 +59,31 @@ def test_sa4():
[2, 3],
[4, 0]])
assert np.all(sa == sa_correct)
+
+
+def test_sa5():
+ """
+ Test wrapping.
+ """
+ a = np.arange(5)
+ sa = segment_axis(a, 5, 0, end='wrap')
+ sa_correct = np.array([[0, 1, 2, 3, 4]])
+ assert np.all(sa == sa_correct)
+
+
+def test_sa6():
+ """
+ Test exceptions.
+ """
+ a = np.arange(5)
+ with pytest.raises(ValueError):
+ sa = segment_axis(a, 2, 3, end='wrap')
+
+
+def test_sa7():
+ """
+ Test exceptions.
+ """
+ a = np.arange(5)
+ with pytest.raises(ValueError):
+ sa = segment_axis(a, 2, -1, end='wrap')
