def test_can_access_sub_dag():
sub_dag = DAG('sub_dag')
ta = PythonCallable(touch_root, File('a.txt'), sub_dag, 'ta')
tb = PythonCallable(touch, File('b.txt'), sub_dag, 'tb')
tc = PythonCallable(touch, File('c.txt'), sub_dag, 'tc')
ta >> tb >> tc
dag = DAG('dag')
fd = Path('d.txt')
td = PythonCallable(touch, File(fd), dag, 'td')
td.set_upstream(sub_dag)
assert 'sub_dag' in td.upstream
def test_task_grouping():
dag = DAG()
t1 = PythonCallable(touch_root, File('1.txt'), dag, name='first')
t2 = PythonCallable(touch_root, File('2.txt'), dag, name='second')
t3 = PythonCallable(touch, File('3.txt'), dag, name='third')
t3.set_upstream(t1, group_name='group')
t3.set_upstream(t2, group_name='group')
dag.render()
assert set(t3.upstream) == {'first', 'second'}
assert set(t3._upstream_product_grouped) == {'group'}
assert set(t3._upstream_product_grouped['group']) == {'first', 'second'}
assert set(t3.params['upstream']) == {'group'}
assert t3.params['upstream']['group']['first'] is t1.product
assert t3.params['upstream']['group']['second'] is t2.product
