class TestGraph(unittest.TestCase):
def setUp(self):
self.test_graph = Graph()
self.test_node_0 = Node(0, 0.1, 1, 1, 1, 1, 1, 1)
self.test_node_1 = Node(1, 0.1, 1, 1, 1, 1, 1, 1)
def test_add_vertex(self):
# add the two test nodes
self.test_graph.add_vertex(self.test_node_0)
self.test_graph.add_vertex(self.test_node_1)
# see if they are in the graph vertex dictionary
self.assertEqual(len(self.test_graph.vertices), 2)
self.assertTrue(0 in self.test_graph.vertices)
self.assertTrue(1 in self.test_graph.vertices)
# if node already in graph, do not allow adding
with self.assertRaises(AssertionError):
self.test_graph.add_vertex(self.test_node_0)
def test_add_edge(self):
n0, n1 = self.test_node_0, self.test_node_1
strength: int = 10
self.test_graph.add_vertex(n0)
self.test_graph.add_vertex(n1)
self.test_graph.add_edge(n0, n1, strength)
self.assertTrue(1 in n0.neighbours)
self.assertTrue(0 in n1.neighbours)
# check bidirectional strength
self.assertEqual(n0.neighbours[1], strength)
self.assertEqual(n1.neighbours[0], strength)
# do not allow adding the same edge
with self.assertRaises(AssertionError):
self.test_graph.add_edge(n0, n1, strength)
with self.assertRaises(AssertionError):
self.test_graph.add_edge(n1, n0, strength)
def test_remove_vertex(self):
n0, n1 = self.test_node_0, self.test_node_1
n2 = Node(2, 0.1, 1, 1, 1, 1, 1, 1)
self.test_graph.add_vertex(n0)
self.test_graph.add_vertex(n1)
self.test_graph.add_vertex(n2)
self.test_graph.add_edge(n0, n1, 5)
self.test_graph.add_edge(n1, n2, 15)
self.test_graph.add_edge(n2, n0, 25)
self.test_graph.remove_vertex(0)
# test_node_0 no longer in neighbours dicts
self.assertFalse(0 in n1.neighbours)
self.assertFalse(0 in n2.neighbours)
# neighbours edges intact
self.assertTrue(2 in n1.neighbours)
self.assertTrue(1 in n2.neighbours)
self.assertFalse(0 in self.test_graph.vertices)
# try to remove something that is removed
with self.assertRaises(AssertionError):
self.test_graph.remove_vertex(0)
self.test_graph.remove_vertex(1)
self.assertFalse(1 in n2.neighbours)
def test_get_node(self):
self.test_graph.add_vertex(self.test_node_0)
self.assertEqual(self.test_graph.get_node(0), self.test_node_0)
with self.assertRaises(AssertionError):
self.test_graph.get_node(1)
