def test_flow_runner_runs_basic_flow_with_1_task():
flow = Flow(name="test")
task = SuccessTask()
flow.add_task(task)
flow_runner = FlowRunner(flow=flow)
state = flow_runner.run(return_tasks=[task])
assert state == Success(result={task: Success(result=1)})
def test_parameters_are_placed_into_context():
flow = Flow(name="test")
y = prefect.Parameter("y", default=99)
flow.add_task(y)
flow_state = FlowRunner(flow=flow).run(return_tasks=[y],
parameters=dict(y=42))
assert isinstance(flow_state, Success)
assert flow_state.result[y].result == 42
def test_required_parameters_must_be_provided():
flow = Flow(name="test")
y = prefect.Parameter("y")
flow.add_task(y)
flow_state = FlowRunner(flow=flow).run(return_tasks=[y])
assert isinstance(flow_state, Failed)
assert isinstance(flow_state.result[y], Failed)
assert "required but not provided" in str(flow_state.result[y]).lower()
def test_flow_runner_with_no_return_tasks():
"""
Make sure FlowRunner accepts return_tasks=None and doesn't raise early error
"""
flow = Flow(name="test")
task = SuccessTask()
flow.add_task(task)
flow_runner = FlowRunner(flow=flow)
assert flow_runner.run(return_tasks=None)
def test_deserialize_with_parameters_key():
f = Flow(name="test")
x = Parameter("x")
f.add_task(x)
f2 = FlowSchema().load(FlowSchema().dump(f))
assert {p.name for p in f2.parameters()} == {p.name for p in f.parameters()}
f_params = {(p.name, p.required, p.default) for p in f.parameters()}
f2_params = {(p.name, p.required, p.default) for p in f2.parameters()}
assert f_params == f2_params
def test_raise_on_exception_plays_well_with_context():
flow = Flow(name="test")
flow.add_task(MathTask())
try:
assert "raise_on_exception" not in prefect.context
with raise_on_exception():
assert "raise_on_exception" in prefect.context
flow.run()
except ZeroDivisionError:
assert "raise_on_exception" not in prefect.context
pass
def test_nested_collection_automatically_applied_to_callargs_imperative(self):
x = Parameter("x")
y = Parameter("y")
identity = IdentityTask()
f = Flow(name="test")
f.add_task(identity)
identity.bind(x=dict(a=[x, dict(y=y)], b=(y, set([x]))), flow=f)
state = f.run(parameters=dict(x=1, y=2))
assert len(f.tasks) == 10
assert state.result[identity].result == dict(a=[1, dict(y=2)], b=(2, set([1])))
def test_set_automatically_applied_to_callargs_imperative(self):
x = Parameter("x")
y = Parameter("y")
identity = IdentityTask()
f = Flow(name="test")
f.add_task(identity)
identity.bind(x=set([x, y]), flow=f)
state = f.run(parameters=dict(x=1, y=2))
assert len(f.tasks) == 4
assert sum(isinstance(t, collections.Set) for t in f.tasks) == 1
assert state.result[identity].result == set([1, 2])
def test_dict_automatically_applied_to_callargs_imperative(self):
x = Parameter("x")
y = Parameter("y")
identity = IdentityTask()
f = Flow(name="test")
f.add_task(identity)
identity.bind(x=dict(a=x, b=y), flow=f)
state = f.run(parameters=dict(x=1, y=2))
assert len(f.tasks) == 5 # 2 params, identity, Dict, List of dict values
assert sum(isinstance(t, collections.Dict) for t in f.tasks) == 1
assert state.result[identity].result == dict(a=1, b=2)
def test_flow_runner_runs_basic_flow_with_2_independent_tasks():
flow = Flow(name="test")
task1 = SuccessTask()
task2 = SuccessTask()
flow.add_task(task1)
flow.add_task(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_run_count_updates_after_each_retry(self):
flow = Flow(name="test")
t1 = ErrorTask(max_retries=2, retry_delay=datetime.timedelta(0))
flow.add_task(t1)
state1 = FlowRunner(flow=flow).run(return_tasks=[t1])
assert state1.result[t1].is_retrying()
assert state1.result[t1].run_count == 1
state2 = FlowRunner(flow=flow).run(return_tasks=[t1], task_states=state1.result)
assert state2.result[t1].is_retrying()
assert state2.result[t1].run_count == 2
def test_parameters():
f = Flow(name="test")
x = Parameter("x")
y = Parameter("y", default=5)
f.add_task(x)
f.add_task(y)
serialized = FlowSchema().dump(f)
assert "parameters" in serialized
assert [
isinstance(ParameterSchema().load(p), Parameter)
for p in serialized["parameters"]
]
def test_raise_on_exception_ignores_all_prefect_signals(signal):
flow = Flow(name="test")
@prefect.task
def raise_signal():
if (prefect.context.get("task_loop_count", 1) < 2
and prefect.context.get("task_run_count", 1) < 2
and signal.__name__ != "PAUSE"):
raise signal("my message")
flow.add_task(raise_signal)
with raise_on_exception():
flow_state = flow.run()
def test_that_bad_code_in_flow_runner_is_caught():
"""
Test that an error in the actual FlowRunner code itself (not the execution of user code)
is caught
"""
class BadFlowRunner(FlowRunner):
def get_flow_run_state(self, *args, **kwargs):
raise RuntimeError("I represent bad code in the flow runner.")
flow = Flow(name="test")
flow.add_task(MathTask())
flow_state = BadFlowRunner(flow=flow).run()
assert isinstance(flow_state, state.Failed)
assert isinstance(flow_state.result, RuntimeError)
assert "I represent bad code in the flow runner." == str(flow_state.result)
def test_secrets_dynamically_pull_from_context():
flow = Flow(name="test")
task1 = PrefectSecret("foo", max_retries=1, retry_delay=datetime.timedelta(0))
flow.add_task(task1)
flow_state = FlowRunner(flow=flow).run(return_tasks=[task1])
assert flow_state.is_running()
assert flow_state.result[task1].is_retrying()
with prefect.context(secrets=dict(foo=42)):
time.sleep(1)
flow_state = FlowRunner(flow=flow).run(task_states=flow_state.result)
assert flow_state.is_successful()
def test_that_bad_code_in_task_runner_is_caught():
"""
Test that an error in the actual TaskRunner code itself (not the execution of user code)
is caught
"""
class BadTaskRunner(TaskRunner):
def get_task_run_state(self, *args, **kwargs):
raise RuntimeError("I represent bad code in the task runner.")
flow = Flow(name="test")
math_task = MathTask()
flow.add_task(math_task)
flow_state = FlowRunner(
flow=flow, task_runner_cls=BadTaskRunner).run(return_tasks=[math_task])
assert isinstance(flow_state, state.Failed)
assert flow_state.result[math_task].is_failed()
caught_exc = flow_state.result[math_task].result
assert isinstance(caught_exc, RuntimeError)
assert "I represent bad code in the task runner." == str(caught_exc)
def test_raise_on_exception_raises_basic_error():
flow = Flow(name="test")
flow.add_task(MathTask())
with pytest.raises(ZeroDivisionError):
with raise_on_exception():
flow.run()
