add unit tests

tests/unit/test_action.py

@@ -6,6 +6,8 @@ import torch
 from rubik.action import (
     POS_ROTATIONS,
     POS_SHIFTS,
+    FACE_ROTATIONS,
+    build_actions_tensor,
 )
 
 
@@ -57,3 +59,41 @@ def test_position_shift(axis: int, size: int, input: Iterable[int], expected: It
     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}"
+
+
+def test_face_rotation_shape():
+    """
+    Test that FACE_ROTATIONS has expected shape.
+    """
+    expected = (3, 6, 6)
+    observed = FACE_ROTATIONS.shape
+    assert expected == observed, f"Face rotation tensor expected shape '{expected}', got '{observed}' instead"
+
+
+@pytest.mark.parametrize(
+    "axis, input, expected",
+    [
+        (0, (1, 0, 0, 0, 0, 0), (0, 0, 1, 0, 0, 0)),  # rotation about X axis: 0 (Up) -> 2 (Front)
+        (1, (1, 0, 0, 0, 0, 0), (0, 1, 0, 0, 0, 0)),  # rotation about Y axis: 0 (Up) -> 1 (Left)
+        (2, (0, 1, 0, 0, 0, 0), (0, 0, 1, 0, 0, 0)),  # rotation about Z axis: 1 (Left) -> 2 (Front)
+    ],
+)
+def test_face_rotation(axis: int, input: Iterable[int], expected: Iterable[int]):
+    """
+    Test that POS_ROTATIONS behaves as expected.
+    """
+    out = torch.tensor(input, dtype=FACE_ROTATIONS.dtype) @ FACE_ROTATIONS[axis]
+    exp = torch.tensor(expected, dtype=FACE_ROTATIONS.dtype)
+    assert torch.equal(out, exp), f"Face rotation tensor is incorrect along axis {axis}: {out} != {exp}"
+
+
+@pytest.mark.parametrize("size", [2, 3, 5, 20])
+def test_build_actions_tensor_shape(size: int):
+    """
+    Test that "build_actions_tensor" output has expected shape.
+    """
+    expected = (3, size, 2, 6 * (size**2), 6 * (size**2))
+    observed = build_actions_tensor(size).shape
+    assert expected == observed, (
+        f"'build_actions_tensor' output has incorrect shape: expected shape '{expected}', got '{observed}' instead"
+    )