Add test for graph initialisation

Author: IsaacCheng9

tests/test_boruvka.py

@@ -1,6 +1,10 @@
 import pytest
 
-from boruvkas_algorithm.boruvka import Graph, find_mst_with_boruvkas_algorithm
+from boruvkas_algorithm.boruvka import (
+    Graph,
+    UnionFind,
+    find_mst_with_boruvkas_algorithm,
+)
 
 
 @pytest.fixture
@@ -15,6 +19,16 @@ def setup_graph():
     return Graph(9)  # Example graph with 9 vertices.
 
 
+def test_graph_initialization():
+    """
+    Tests that a graph is initialised with the correct number of vertices and
+    no edges.
+    """
+    graph = Graph(5)
+    assert len(graph.vertices) == 5, "Graph should have 5 vertices"
+    assert len(graph.edges) == 0, "Graph should be initialised with no edges"
+
+
 def test_add_edge(setup_graph: Graph):
     """
     Tests that edges are correctly added by checking the length of the edge
@@ -80,11 +94,54 @@ def test_mst(setup_graph: Graph):
     )
 
 
-def test_graph_initialization():
-    """
-    Test that a graph is initialized with the correct number of vertices and
-    no edges.
-    """
-    graph = Graph(5)  # Initialize a graph with 5 vertices.
-    assert len(graph.vertices) == 5, "Graph should have 5 vertices"
-    assert len(graph.edges) == 0, "Graph should be initialized with no edges"
+def test_mst_with_injected_union_find(setup_graph: Graph):
+    """Tests that the algorithm works with an injected UnionFind instance."""
+    graph = setup_graph
+    graph.add_edge(0, 1, 4)
+    graph.add_edge(0, 6, 7)
+    graph.add_edge(1, 6, 11)
+    graph.add_edge(1, 7, 20)
+    graph.add_edge(1, 2, 9)
+    graph.add_edge(2, 3, 6)
+    graph.add_edge(2, 4, 2)
+    graph.add_edge(3, 4, 10)
+    graph.add_edge(3, 5, 5)
+    graph.add_edge(4, 5, 15)
+    graph.add_edge(4, 7, 1)
+    graph.add_edge(4, 8, 5)
+    graph.add_edge(5, 8, 12)
+    graph.add_edge(6, 7, 1)
+    graph.add_edge(7, 8, 3)
+
+    # Inject a custom UnionFind instance.
+    union_find = UnionFind(len(graph.vertices))
+    mst_weight, mst_edges = find_mst_with_boruvkas_algorithm(graph, union_find)
+
+    assert mst_weight == 29, "MST weight should be 29"
+    assert len(mst_edges) == 8, "MST should have 8 edges for 9 vertices"
+
+
+def test_mst_simple_triangle():
+    """Tests MST on a simple triangle graph."""
+    graph = Graph(3)
+    graph.add_edge(0, 1, 1)
+    graph.add_edge(1, 2, 2)
+    graph.add_edge(0, 2, 3)
+
+    mst_weight, mst_edges = find_mst_with_boruvkas_algorithm(graph)
+
+    assert mst_weight == 3, "MST weight should be 3 (edges 1 + 2)"
+    assert len(mst_edges) == 2, "MST should have 2 edges for 3 vertices"
+
+
+def test_mst_linear_graph():
+    """Tests MST on a linear graph (already a tree)."""
+    graph = Graph(4)
+    graph.add_edge(0, 1, 1)
+    graph.add_edge(1, 2, 2)
+    graph.add_edge(2, 3, 3)
+
+    mst_weight, mst_edges = find_mst_with_boruvkas_algorithm(graph)
+
+    assert mst_weight == 6, "MST weight should be 6 (1 + 2 + 3)"
+    assert len(mst_edges) == 3, "MST should have 3 edges for 4 vertices"