Added tests to tests/test_matrix and modified __init__.py

keon · Oct 2, 2020 · d1e1956 · d1e1956
1 parent 0e61ad8
commit d1e1956
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 20 deletions.
diff --git a/algorithms/matrix/__init__.py b/algorithms/matrix/__init__.py
@@ -1 +0,0 @@
-from .matrix_exponentiation import *

diff --git a/algorithms/matrix/matrix_exponentiation.py b/algorithms/matrix/matrix_exponentiation.py
@@ -39,22 +39,3 @@ def matrix_exponentiation(mat: list, n: int) -> list:
     else:
         tmp = matrix_exponentiation(mat, n // 2)
         return multiply(tmp, tmp)
-
-if __name__ == "__main__":
-    mat = [[1, 0, 2], [2, 1, 0], [0, 2, 1]]
-
-    res0 = matrix_exponentiation(mat, 0)
-    assert res0 == [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
-    print(f"{mat}^0 = {res0}")
-
-    res1 = matrix_exponentiation(mat, 1)
-    assert res1 == [[1, 0, 2], [2, 1, 0], [0, 2, 1]]
-    print(f"{mat}^1 = {res1}")
-
-    res2 = matrix_exponentiation(mat, 2)
-    assert res2 == [[1, 4, 4], [4, 1, 4], [4, 4, 1]]
-    print(f"{mat}^2 = {res2}")
-
-    res5 = matrix_exponentiation(mat, 5)
-    assert res5 == [[81, 72, 90], [90, 81, 72], [72, 90, 81]]
-    print(f"{mat}^5 = {res5}")
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
@@ -3,6 +3,7 @@
     copy_transform,
     crout_matrix_decomposition,
     cholesky_matrix_decomposition,
+    matrix_exponentiation,
     matrix_inversion,
     multiply,
     rotate_image,
@@ -159,6 +160,29 @@ def test_inversion(self):
                                                               [Fraction(13, 53), Fraction(-22, 53), Fraction(5, 53)]])
 
 
+class TestMatrixExponentiation(unittest.TestCase):
+    """[summary]
+    Test for the file matrix_exponentiation.py
+
+    Arguments:
+        unittest {[type]} -- [description]
+    """
+
+    def test_matrix_exponentiation(self):
+        mat = [[1, 0, 2], [2, 1, 0], [0, 2, 1]]
+
+        self.assertEqual(matrix_exponentiation.matrix_exponentiation(mat, 0),
+            [[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+
+        self.assertEqual(matrix_exponentiation.matrix_exponentiation(mat, 1),
+            [[1, 0, 2], [2, 1, 0], [0, 2, 1]])
+
+        self.assertEqual(matrix_exponentiation.matrix_exponentiation(mat, 2),
+            [[1, 4, 4], [4, 1, 4], [4, 4, 1]])
+
+        self.assertEqual(matrix_exponentiation.matrix_exponentiation(mat, 5),
+            [[81, 72, 90], [90, 81, 72], [72, 90, 81]])
+
 
 class TestMultiply(unittest.TestCase):
     """[summary]