def test_runs_on_finish(executor, tmp_directory):
hook.count = 0
hook_2.count = 0
hook_3.count = 0
hook_4.count = 0
dag = DAG(executor=executor)
t = PythonCallable(fn, File('file1.txt'), dag, 't')
t.on_finish = hook
t.on_failure = hook_4
t2 = PythonCallable(touch_w_upstream, File('file2'), dag, 't2')
t2.on_finish = hook_2
t3 = PythonCallable(fn, File('file3'), dag, 't3')
t3.on_finish = hook_3
t >> t2
dag.build()
assert hook.count == 1
assert hook_2.count == 1
assert hook_3.count == 1
assert hook_4.count == 0
def test_on_finish(tmp_directory):
dag = DAG()
t = PythonCallable(touch, File('file'), dag, name='touch')
t.on_finish = on_finish
dag.build()
def test_runs_on_finish(tmp_directory, capsys):
dag = DAG()
t = PythonCallable(fn1, File('file1.txt'), dag, name='fn1')
t.on_finish = on_finish
dag.build()
assert capsys.readouterr().out == 'running on finish\n'
def test_on_finish_exceptions_are_logged(executor, tmp_directory, caplog):
dag = DAG(executor=executor)
t = PythonCallable(fn, File('file.txt'), dag, name='t')
t.on_finish = hook_crashing
with caplog.at_level(logging.ERROR):
with pytest.raises(DAGBuildError):
dag.build()
assert 'Exception when running on_finish for task "t"' in caplog.text
def test_task_status_and_output_when_on_finish_crashes(tmp_directory):
dag = DAG()
t = PythonCallable(fn, File('file'), dag)
t.on_finish = hook_crashing
t2 = PythonCallable(touch_w_upstream, File('file2'), dag)
t >> t2
with pytest.raises(DAGBuildError) as excinfo:
dag.build()
assert t.exec_status == TaskStatus.Errored
assert t2.exec_status == TaskStatus.Aborted
assert "PythonCallable: fn -> File('file')" in str(excinfo.getrepr())
def test_early_stop_from_task_level_on_finish(executor, tmp_directory):
dag = DAG(executor=executor)
t = PythonCallable(touch_root, File('file.txt'), dag)
t.on_finish = early_stop
assert dag.build() is None
def test_passes_client_to_on_finish(tmp_directory):
dag = DAG()
t = PythonCallable(fn1, File('file1.txt'), dag, name='fn1')
t.on_finish = on_finish_w_client
dag.build()
"""
df = pd.read_csv(str(upstream['dump']))
df['a'] = df['a'] + 1
df.to_csv(str(product), index=False)
def check_a_has_no_nas(task):
df = pd.read_csv(str(task.product))
# this print is just here to show that the hook is executed
print('\n\nRunning on_finish hook: checking column has no NAs...\n\n')
assert not df.a.isna().sum()
task_add_one = PythonCallable(_add_one,
File(tmp_dir / 'add_one.csv'),
dag,
name='add_one')
task_add_one.on_finish = check_a_has_no_nas
task_dump >> task_add_one
###############################################################################
# Pipeline plot
# -------------
dag.plot(output='matplotlib')
###############################################################################
# Pipeline build
# -------------
dag.build()
def make():
# NOTE: must run callables in the same process so counting works
dag = DAG(executor=Serial(build_in_subprocess=False))
t = PythonCallable(fn, File('file1.txt'), dag)
t.on_finish = hook_crashing
return dag
