def test_flow_runner_doesnt_return_by_default():
flow = Flow(name="test")
task1 = SuccessTask()
task2 = SuccessTask()
flow.add_edge(task1, task2)
res = FlowRunner(flow=flow).run()
assert res.result == {}
def test_new_edge_objects_can_test_membership_in_flow():
flow = Flow(name="test")
t1 = TaskWithKey()
t2 = TaskWithKey()
flow.add_edge(t1, t2, key="a_key")
assert Edge(t1, t2, key="a_key") in flow.edges
def test_flow_with_multiple_retry_tasks_doesnt_run_them_early():
"""
t1 -> t2
t1 -> t3
Both t2 and t3 fail on initial run and request a retry. Starting the flow at t1 should
only run t3, which requests an immediate retry, and not t2, which requests a retry in
10 minutes.
This tests a check on the TaskRunner, but which matters in Flows like this.
"""
flow = Flow(name="test")
t1 = Task()
t2 = ErrorTask(retry_delay=datetime.timedelta(minutes=10), max_retries=1)
t3 = ErrorTask(retry_delay=datetime.timedelta(minutes=0), max_retries=1)
flow.add_edge(t1, t2)
flow.add_edge(t1, t3)
state1 = FlowRunner(flow=flow).run(return_tasks=flow.tasks)
assert state1.result[t2].is_retrying()
assert state1.result[t3].is_retrying()
state2 = FlowRunner(flow=flow).run(return_tasks=flow.tasks,
task_states=state1.result)
assert state2.result[t2].is_retrying()
assert state2.result[t2] == state1.result[
t2] # state is not modified at all
assert isinstance(state2.result[t3], Failed) # this task ran
def test_flow_runner_does_return_tasks_when_requested():
flow = Flow(name="test")
task1 = SuccessTask()
task2 = SuccessTask()
flow.add_edge(task1, task2)
flow_state = FlowRunner(flow=flow).run(return_tasks=[task1])
assert isinstance(flow_state, Success)
assert isinstance(flow_state.result[task1], Success)
def test_flow_runner_runs_base_task_class():
flow = Flow(name="test")
task1 = Task()
task2 = Task()
flow.add_edge(task1, task2)
flow_state = FlowRunner(flow=flow).run(return_tasks=[task1, task2])
assert isinstance(flow_state, Success)
assert isinstance(flow_state.result[task1], Success)
assert isinstance(flow_state.result[task2], Success)
def test_flow_run_state_determined_by_reference_tasks():
flow = Flow(name="test")
t1 = ErrorTask()
t2 = SuccessTask(trigger=prefect.triggers.all_finished)
flow.add_edge(t1, t2)
flow.set_reference_tasks([t1])
flow_state = flow.run()
assert isinstance(flow_state, Failed)
assert isinstance(flow_state.result[t1], Failed)
assert isinstance(flow_state.result[t2], Success)
def test_flow_runner_does_not_return_task_states_when_it_doesnt_run():
flow = Flow(name="test")
task1 = SuccessTask()
task2 = ErrorTask()
flow.add_edge(task1, task2)
flow_state = FlowRunner(flow=flow).run(state=Success(result=5),
return_tasks=[task1, task2])
assert isinstance(flow_state, Success)
assert flow_state.result == 5
def test_flow_runner_runs_basic_flow_with_2_dependent_tasks():
flow = Flow(name="test")
task1 = SuccessTask()
task2 = SuccessTask()
flow.add_edge(task1, task2)
flow_state = FlowRunner(flow=flow).run(return_tasks=[task1, task2])
assert isinstance(flow_state, Success)
assert flow_state.result[task1] == Success(result=1)
assert flow_state.result[task2] == Success(result=1)
def test_flow_run_state_not_determined_by_reference_tasks_if_terminal_tasks_are_not_finished():
flow = Flow(name="test")
t1 = ErrorTask()
t2 = RaiseRetryTask(trigger=prefect.triggers.all_finished)
flow.add_edge(t1, t2)
flow.set_reference_tasks([t1])
flow_state = FlowRunner(flow=flow).run(return_tasks=flow.tasks)
assert flow_state.is_running()
assert flow_state.result[t1].is_failed()
assert flow_state.result[t2].is_retrying()
def test_flow_runner_runs_basic_flow_with_2_dependent_tasks_and_second_task_fails():
flow = Flow(name="test")
task1 = SuccessTask()
task2 = ErrorTask()
flow.add_edge(task1, task2)
flow_state = FlowRunner(flow=flow).run(return_tasks=[task1, task2])
assert isinstance(flow_state, Failed)
assert isinstance(flow_state.result[task1], Success)
assert isinstance(flow_state.result[task2], Failed)
def test_flow_runner_makes_copy_of_task_results_dict():
"""
Ensure the flow runner copies the task_results dict rather than modifying it inplace
"""
flow = Flow(name="test")
t1, t2 = Task(), Task()
flow.add_edge(t1, t2)
task_states = {t1: Pending()}
state = flow.run(task_states=task_states)
assert state.result[t1] == Success(result=None)
assert task_states == {t1: Pending()}
def test_flow_runner_remains_running_if_tasks_are_retrying():
# https://github.com/PrefectHQ/prefect/issues/19
flow = Flow(name="test")
task1 = SuccessTask()
task2 = ErrorTask(max_retries=1, retry_delay=datetime.timedelta(0))
flow.add_edge(task1, task2)
flow_state = FlowRunner(flow=flow).run(return_tasks=[task1, task2])
assert flow_state.is_running()
assert flow_state.result[task1].is_successful()
assert flow_state.result[task2].is_retrying()
def test_flow_run_method_returns_task_states_even_if_it_doesnt_run():
# https://github.com/PrefectHQ/prefect/issues/19
flow = Flow(name="test")
task1 = SuccessTask()
task2 = ErrorTask()
flow.add_edge(task1, task2)
flow_state = flow.run(state=Success())
assert flow_state.is_successful()
assert flow_state.result[task1].is_pending()
assert flow_state.result[task2].is_pending()
def test_flow_runner_runs_basic_flow_with_2_dependent_tasks_and_first_task_fails_with_FAIL():
flow = Flow(name="test")
task1 = RaiseFailTask()
task2 = SuccessTask()
flow.add_edge(task1, task2)
flow_state = FlowRunner(flow=flow).run(return_tasks=[task1, task2])
assert isinstance(flow_state, Failed)
assert isinstance(flow_state.result[task1], Failed)
assert not isinstance(flow_state.result[task1], TriggerFailed)
assert isinstance(flow_state.result[task2], TriggerFailed)
def test_flow_runner_runs_flow_with_2_dependent_tasks_and_first_task_fails_and_second_has_trigger(
):
flow = Flow(name="test")
task1 = ErrorTask()
task2 = SuccessTask(trigger=prefect.triggers.all_failed)
flow.add_edge(task1, task2)
flow_state = FlowRunner(flow=flow).run(return_tasks=[task1, task2])
assert isinstance(flow_state,
Success) # flow state is determined by terminal states
assert isinstance(flow_state.result[task1], Failed)
assert isinstance(flow_state.result[task2], Success)
def test_deserialize_edges():
"""
Tests that edges are appropriately deserialized, even in they involve keys.
Also tests that tasks are deserialized in a way that reuses them in edges -- in other
words, when edges are loaded they use their corresponding task IDs to access the
correct Task objects out of a cache.
"""
class ArgTask(Task):
def run(self, x):
return x
f = Flow(name="test")
t1, t2, t3 = Task("a"), Task("b"), ArgTask("c")
f.add_edge(t1, t2)
f.add_edge(t2, t3, key="x")
f.add_edge(t1, t3, mapped=True)
serialized = FlowSchema().dump(f)
deserialized = FlowSchema().load(serialized)
d1, d2, d3 = sorted(deserialized.tasks, key=lambda t: t.name)
assert deserialized.edges == {
Edge(d1, d2),
Edge(d2, d3, key="x"),
Edge(d1, d3, mapped=True),
}
