def test_add_update(self): g = DAG() g.add_or_update_node(1) g.add_or_update_node(2) g.add_or_update_node(3, inputs = {1,2}) g.add_or_update_node(0, outputs={1}) self.assertEqual(g._head_nodes, g.to_nodes({0,2})) self.assertEqual(g._leaf_nodes, g.to_nodes({3})) g.add_or_update_node(0, outputs={1,2,3}) with self.assertRaises(CycleError): g.add_or_update_node(0, inputs={1}) self.assertEqual(g[0].outputs, g.to_nodes({1,2,3}))
def test_complex(self): spec = ( (1, (2,3,4,5)), (2, ()), (3, ()), (4, ()), (5, (2,)), ) g = DAG() for val, deps in spec: for dep in deps: g.add_or_update_node(dep) g.add_or_update_node(val, inputs=deps) self.assertEqual(g._head_nodes, g.to_nodes({2,3,4})) g = DAG() for val, deps in spec: for dep in deps: g.add_or_update_node(dep) g.add_or_update_node(val, outputs=deps) self.assertEqual(g._leaf_nodes, g.to_nodes({2,3,4}))