diff --git a/.idea/matrixmania.iml b/.idea/matrixmania.iml
index 54daf84..d42f524 100644
--- a/.idea/matrixmania.iml
+++ b/.idea/matrixmania.iml
@@ -1,7 +1,9 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <module type="PYTHON_MODULE" version="4">
   <component name="NewModuleRootManager">
-    <content url="file://$MODULE_DIR$" />
+    <content url="file://$MODULE_DIR$">
+      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+    </content>
     <orderEntry type="jdk" jdkName="Python 3.13 (PythonProject)" jdkType="Python SDK" />
     <orderEntry type="sourceFolder" forTests="false" />
   </component>
diff --git a/test/test_matrixmania/test_rot.py b/test/test_matrixmania/test_rot.py
index 3eba5c5..5555db0 100644
--- a/test/test_matrixmania/test_rot.py
+++ b/test/test_matrixmania/test_rot.py
@@ -2,7 +2,7 @@ import math
 
 import pytest
 
-from matrixmania.compute import matmul, rot_2D, rot_3D, transpose
+from matrixmania.compute import rot_2D, rot_3D
 
 
 def assert_matrix_close(actual, expected):
@@ -14,45 +14,41 @@ def assert_matrix_close(actual, expected):
 
 # ---------- rot_2D ----------
 
-def test_rot_2D_zero():
-    assert_matrix_close(rot_2D(0), [[1, 0], [0, 1]])
-
-
 def test_rot_2D_90_degree():
-    assert_matrix_close(rot_2D(math.pi / 2), [[0, -1], [1, 0]])
+    # Arrange
+    angle = math.pi / 2
+    expected = [[0, -1],
+                [1, 0]]
 
+    # Act
+    result = rot_2D(angle)
 
-def test_rot_2D_invalid_input():
-    with pytest.raises(TypeError):
-        rot_2D("abc")
+    # Assert
+    assert_matrix_close(result, expected)
 
 
 # ---------- rot_3D ----------
 
-@pytest.mark.parametrize("axis", ["x", "y", "z"])
-def test_rot_3D_zero_is_identity(axis):
-    assert_matrix_close(rot_3D(0, axis), [[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+def test_rot_3D_90_degree_z():
+    # Arrange
+    angle = math.pi / 2
+    axis = "z"
+    expected = [[0, -1, 0],
+                [1, 0, 0],
+                [0, 0, 1]]
+
+    # Act
+    result = rot_3D(angle, axis)
+
+    # Assert
+    assert_matrix_close(result, expected)
 
 
-@pytest.mark.parametrize(
-    "axis, expected",
-    [
-        ("x", [[1, 0, 0], [0, 0, -1], [0, 1, 0]]),
-        ("y", [[0, 0, 1], [0, 1, 0], [-1, 0, 0]]),
-        ("z", [[0, -1, 0], [1, 0, 0], [0, 0, 1]]),
-    ],
-)
-def test_rot_3D_90_degree(axis, expected):
-    assert_matrix_close(rot_3D(math.pi / 2, axis), expected)
+def test_rot_3D_invalid_axis():
+    # Arrange
+    angle = 1.0
+    axis = "w"
 
-
-@pytest.mark.parametrize("axis", ["x", "y", "z"])
-def test_rot_3D_is_orthogonal(axis):
-    r = rot_3D(1.0, axis)
-    assert_matrix_close(matmul(r, transpose(r)), [[1, 0, 0], [0, 1, 0], [0, 0, 1]])
-
-
-@pytest.mark.parametrize("axis", ["w", "X", "", "xy"])
-def test_rot_3D_invalid_axis(axis):
+    # Act + Assert
     with pytest.raises(ValueError):
-        rot_3D(1.0, axis)
\ No newline at end of file
+        rot_3D(angle, axis)
\ No newline at end of file