def test_remove_node(self): """ Tests removing a node from the graph. """ g = DAG() g.add_node(1) g.add_node(2) g.add_node(3) g.add_edge(1, 2) g.add_edge(1, 3) g.add_edge(2, 3) g.remove_node(3) self.assertNotIn(3, g.all_nodes) self.assertEqual(len(g.dependent_nodes(1)), 1) self.assertIn(2, g.dependent_nodes(1)) self.assertEqual(len(g.dependent_nodes(2)), 0) g.remove_node(1) self.assertEqual(g.in_degree[g.nodes_map[2].id], 0)
def test_add_edge(self): """ Tests adding edges to a DAG. """ g = DAG() g.add_node(1) g.add_node(2) g.add_edge(1, 2) self.assertEqual(len(g.dependent_nodes(1)), 1) self.assertIn(2, g.dependent_nodes(1)) # In-degrees updated self.assertEqual(g.in_degree[g.nodes_map[1].id], 0) self.assertEqual(g.in_degree[g.nodes_map[2].id], 1) g.add_node(3) g.add_edge(1, 3) self.assertEqual(len(g.dependent_nodes(1)), 2) self.assertIn(3, g.dependent_nodes(1)) # In-degrees updated self.assertEqual(g.in_degree[g.nodes_map[1].id], 0) self.assertEqual(g.in_degree[g.nodes_map[3].id], 1) g.add_edge(2, 3) self.assertEqual(len(g.dependent_nodes(2)), 1) self.assertIn(3, g.dependent_nodes(2)) # In-degrees updated self.assertEqual(g.in_degree[g.nodes_map[3].id], 2) # Add a same edge again - nothing changed? g.add_edge(1, 3) self.assertEqual(len(g.dependent_nodes(1)), 2) # Add an edge resulting in a cycle with self.assertRaises(InvalidDAGException): g.add_edge(3, 1)