add unit test

tests/unit/test_action.py

@@ -5,6 +5,7 @@ import torch
 
 from rubik.action import (
     POS_ROTATIONS,
+    POS_SHIFTS,
 )
 
 
@@ -29,6 +30,27 @@ def test_position_rotation_shape():
     ],
 )
 def test_position_rotation(axis: int, input: Iterable[int], expected: Iterable[int]):
+    """
+    Test that POS_ROTATIONS behaves as expected.
+    """
     out = POS_ROTATIONS[axis] @ torch.tensor(input, dtype=POS_ROTATIONS.dtype)
     exp = torch.tensor(expected, dtype=POS_ROTATIONS.dtype)
     assert torch.equal(out, exp), f"Position rotation tensor is incorrect along axis {axis}: {out} != {exp}"
+
+
+@pytest.mark.parametrize(
+    "axis, size, input, expected",
+    [
+        (0, 3, (1, 1, 1, 1), (1, 1, 1, 0)),
+        (1, 3, (1, 1, 1, 1), (1, 0, 1, 1)),
+        (2, 3, (1, 1, 1, 1), (1, 1, 0, 1)),
+    ],
+)
+def test_position_shift(axis: int, size: int, input: Iterable[int], expected: Iterable[int]):
+    """
+    Test that POS_SHIFTS behaves as expected.
+    """
+    rot = POS_ROTATIONS[axis] @ (torch.tensor(input, dtype=POS_ROTATIONS.dtype) * (size - 1))
+    out = rot + (POS_SHIFTS[axis] * (size - 1))
+    exp = torch.tensor(expected, dtype=POS_ROTATIONS.dtype) * (size - 1)
+    assert torch.equal(out, exp), f"Position shift tensor is incorrect along axis {axis}: {out} != {exp}"