test/graphs/prim: writing a test case to verify the correctness of the algorithm (#2454)

graphs/minimum_spanning_tree_prims.py

@@ -104,7 +104,7 @@ def PrimsAlgorithm(l):  # noqa: E741
     return TreeEdges
 
 
-if __name__ == "__main__":
+if __name__ == "__main__":  # pragma: no cover
     # < --------- Prims Algorithm --------- >
     n = int(input("Enter number of vertices: ").strip())
     e = int(input("Enter number of edges: ").strip())

graphs/tests/test_min_spanning_tree_prim.py

@@ -0,0 +1,47 @@
+from collections import defaultdict
+
+
+from graphs.minimum_spanning_tree_prims import PrimsAlgorithm as mst
+
+
+def test_prim_successful_result():
+    num_nodes, num_edges = 9, 14  # noqa: F841
+    edges = [
+        [0, 1, 4],
+        [0, 7, 8],
+        [1, 2, 8],
+        [7, 8, 7],
+        [7, 6, 1],
+        [2, 8, 2],
+        [8, 6, 6],
+        [2, 3, 7],
+        [2, 5, 4],
+        [6, 5, 2],
+        [3, 5, 14],
+        [3, 4, 9],
+        [5, 4, 10],
+        [1, 7, 11],
+    ]
+
+    adjancency = defaultdict(list)
+    for node1, node2, cost in edges:
+        adjancency[node1].append([node2, cost])
+        adjancency[node2].append([node1, cost])
+
+    result = mst(adjancency)
+
+    expected = [
+        [7, 6, 1],
+        [2, 8, 2],
+        [6, 5, 2],
+        [0, 1, 4],
+        [2, 5, 4],
+        [2, 3, 7],
+        [0, 7, 8],
+        [3, 4, 9],
+    ]
+
+    for answer in expected:
+        edge = tuple(answer[:2])
+        reverse = tuple(edge[::-1])
+        assert edge in result or reverse in result