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