def test_context_includes_date(self):
@prefect.task
def return_ctx_key():
return prefect.context.get("date")
f = Flow(name="test", tasks=[return_ctx_key])
res = f.run()
assert res.is_successful()
output = res.result[return_ctx_key].result
assert isinstance(output, datetime.datetime)
def test_user_provided_context_is_prioritized(self):
@prefect.task
def return_ctx_key():
return prefect.context.get("date")
f = Flow(name="test", tasks=[return_ctx_key])
res = f.run(context={"date": "42"})
assert res.is_successful()
output = res.result[return_ctx_key].result
assert output == "42"
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_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_doesnt_override_scheduled_start_time_when_running_on_schedule(
self, run_on_schedule):
@prefect.task
def return_scheduled_start_time():
return prefect.context.get("scheduled_start_time")
f = Flow(name="test", tasks=[return_scheduled_start_time])
res = f.run(context=dict(scheduled_start_time=42),
run_on_schedule=run_on_schedule)
assert res.is_successful()
assert res.result[return_scheduled_start_time].result != 42
def test_flow_runner_provides_scheduled_start_time(self):
@prefect.task
def return_scheduled_start_time():
return prefect.context.get("scheduled_start_time")
f = Flow(name="test", tasks=[return_scheduled_start_time])
res = f.run()
assert res.is_successful()
assert res.result[return_scheduled_start_time].is_successful()
assert isinstance(res.result[return_scheduled_start_time].result,
datetime.datetime)
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_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_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_context_contains_date_formats(self, date):
@prefect.task
def return_ctx_key():
return prefect.context.get(date)
f = Flow(name="test", tasks=[return_ctx_key])
res = f.run()
assert res.is_successful()
output = res.result[return_ctx_key].result
assert isinstance(output, str)
assert len(output) == 10
def test_flow_runner_captures_and_exposes_dask_errors(executor):
q = queue.Queue()
@prefect.task
def put():
q.put(55)
f = Flow(name="test", tasks=[put])
state = f.run(executor=executor)
assert state.is_failed()
assert isinstance(state.result, TypeError)
assert str(state.result) == "can't pickle _thread.lock objects"
def test_user_provided_context_is_prioritized(self, outer_context,
inner_context, sol):
@prefect.task
def return_ctx_key():
return prefect.context.get("date")
f = Flow(name="test", tasks=[return_ctx_key])
with prefect.context(**outer_context):
res = f.run(context=inner_context)
assert res.is_successful()
output = res.result[return_ctx_key].result
assert output == sol
def test_flow_runner_captures_and_exposes_dask_errors(executor):
q = queue.Queue()
@prefect.task
def put():
q.put(55)
f = Flow(name="test", tasks=[put])
state = f.run(executor=executor)
assert state.is_failed()
assert isinstance(state.result, TypeError)
# assert two possible result outputs for different Python versions
assert str(state.result) in [
"can't pickle _thread.lock objects",
"cannot pickle '_thread.lock' object",
]
def test_flow_runner_properly_provides_context_to_task_runners(executor):
@prefect.task
def my_name():
return prefect.context.get("my_name")
@prefect.task
def flow_name():
return prefect.context.get("flow_name")
flow = Flow(name="test-dummy", tasks=[flow_name, my_name])
with prefect.context(my_name="marvin"):
res = flow.run(executor=executor)
assert res.result[flow_name].result == "test-dummy"
assert res.result[my_name].result == "marvin"
with Flow("test-map") as f:
tt = flow_name.map(upstream_tasks=[my_name])
with prefect.context(my_name="mapped-marvin"):
res = f.run(executor=executor)
assert res.result[my_name].result == "mapped-marvin"
assert res.result[tt].result[0] == "test-map"
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()
def test_task_runners_submitted_to_remote_machines_respect_original_config(
monkeypatch):
"""
This test is meant to simulate the behavior of running a Cloud Flow against an external
cluster which has _not_ been configured for Prefect. The idea is that the configuration
settings which were present on the original machine are respected in the remote job, reflected
here by having the CloudHandler called during logging and the special values present in context.
"""
from prefect.engine.flow_runner import run_task
def my_run_task(*args, **kwargs):
with prefect.utilities.configuration.set_temporary_config({
"logging.log_to_cloud":
False,
"cloud.auth_token":
""
}):
return run_task(*args, **kwargs)
calls = []
class Client:
def write_run_logs(self, *args, **kwargs):
calls.append(args)
monkeypatch.setattr("prefect.engine.flow_runner.run_task", my_run_task)
monkeypatch.setattr("prefect.client.Client", Client)
@prefect.task
def log_stuff():
logger = prefect.context.get("logger")
logger.critical("important log right here")
return (
prefect.context.config.special_key,
prefect.context.config.cloud.auth_token,
)
with prefect.utilities.configuration.set_temporary_config({
"logging.log_to_cloud":
True,
"special_key":
42,
"cloud.auth_token":
"original",
}):
# captures config at init
flow = Flow("test", tasks=[log_stuff])
flow_state = flow.run(task_contexts={log_stuff: dict(special_key=99)})
assert flow_state.is_successful()
assert flow_state.result[log_stuff].result == (42, "original")
time.sleep(0.75)
assert len(calls) >= 1
assert len([log for call in calls
for log in call[0]]) == 5 # actual number of logs
loggers = [log["name"] for call in calls for log in call[0]]
assert set(loggers) == {
"prefect.TaskRunner",
"prefect.FlowRunner",
"prefect.log_stuff",
}
