Unit tests

graphs/graphs_floyd_warshall.py

@@ -6,13 +6,14 @@
 
 
 def _print_dist(dist, v):
-    print("\nThe shortest path matrix using Floyd Warshall algorithm\n")
+    print("The shortest path matrix using Floyd Warshall algorithm\n")
     for i in range(v):
         for j in range(v):
+            end_char = "" if j == v - 1 else "  "
             if dist[i][j] != float("inf"):
-                print(int(dist[i][j]), end="\t")
+                print(int(dist[i][j]), end=end_char)
             else:
-                print("INF", end="\t")
+                print("INF", end=end_char)
         print()
 
 
@@ -31,6 +32,27 @@ def floyd_warshall(graph, v):
     4. The above is repeated for each vertex k in the graph.
     5. Whenever distance[i][j] is given a new minimum value, next vertex[i][j] is
         updated to the next vertex[i][k].
+
+
+    >>> graph = [
+    ...     [0, 3, float('inf')],
+    ...     [2, 0, float('inf')],
+    ...     [float('inf'), 7, 0]
+    ... ]
+
+    >>> expected = [
+    ...     [0, 3, float('inf')],
+    ...     [2, 0, float('inf')],
+    ...     [9, 7, 0]
+    ... ]
+    >>> dist, _ = floyd_warshall(graph, 3)
+    The shortest path matrix using Floyd Warshall algorithm
+    <BLANKLINE>
+    0  3  INF
+    2  0  INF
+    9  7  0
+    >>> dist == expected
+    True
     """
 
     dist = [[float("inf") for _ in range(v)] for _ in range(v)]

graphs/tests/test_graphs_floyd_warshall.py

@@ -0,0 +1,40 @@
+import pytest
+
+from graphs.graphs_floyd_warshall import floyd_warshall
+
+
+def test_no_edges():
+    graph = [
+        [0, float("inf"), float("inf")],
+        [float("inf"), 0, float("inf")],
+        [float("inf"), float("inf"), 0],
+    ]
+    expected = [
+        [0, float("inf"), float("inf")],
+        [float("inf"), 0, float("inf")],
+        [float("inf"), float("inf"), 0],
+    ]
+    dist, _ = floyd_warshall(graph, 3)
+    assert dist == expected
+
+
+def test_with_edges():
+    graph = [[0, 3, float("inf")], [2, 0, float("inf")], [float("inf"), 7, 0]]
+    expected = [[0, 3, float("inf")], [2, 0, float("inf")], [9, 7, 0]]
+    dist, _ = floyd_warshall(graph, 3)
+    assert dist == expected
+
+
+def test_unreachable_vertices():
+    graph = [
+        [0, 1, float("inf")],
+        [float("inf"), 0, 2],
+        [float("inf"), float("inf"), 0],
+    ]
+    expected = [[0, 1, 3], [float("inf"), 0, 2], [float("inf"), float("inf"), 0]]
+    dist, _ = floyd_warshall(graph, 3)
+    assert dist == expected
+
+
+if __name__ == "__main__":
+    pytest.main()