def test_parameters_can_be_set_in_context_if_none_passed():
x = prefect.Parameter("x")
f = FlowRunner(Flow(name="test", tasks=[x]))
state = f.run(parameters={}, context={"parameters": {"x": 5}}, return_tasks=[x])
assert state.result[x].result == 5
def test_set_upstream(self):
f = Flow(name="test")
t1 = Task()
t2 = Task()
t2.set_upstream(t1, flow=f)
assert Edge(t1, t2) in f.edges
def test_set_upstream_no_flow(self):
f = Flow(name="test")
t1 = Task()
t2 = Task()
with pytest.raises(ValueError, match="No Flow was passed"):
t2.set_upstream(t1)
def test_task_runner_handlers_are_called(self):
FlowRunner(flow=Flow(name="test"),
state_handlers=[flow_runner_handler]).run()
# the flow changed state twice: Pending -> Running -> Success
assert handler_results["FlowRunner"] == 2
def test_inputs_task_decorator(self):
with Flow("test"):
assert self.mult(x=1).inputs() == dict(
x=dict(type=Any, required=True, default=None),
y=dict(type=int, required=False, default=1),
)
def test_running_stays_running(self):
state = Running()
flow = Flow(name="test", tasks=[Task()])
new_state = FlowRunner(flow=flow).set_flow_to_running(state=state)
assert new_state.is_running()
def test_flow_on_failure_is_not_called(self):
on_failure = MagicMock()
flow = Flow(name="test", on_failure=on_failure, tasks=[Task()])
FlowRunner(flow=flow).run()
assert not on_failure.called
def test_dict_handles_non_string_keys(self):
l = collections.Dict()
with Flow(name="test") as f:
l.bind(keys=[None, 55], values=[1, 2])
assert f.run().result[l].result == {None: 1, 55: 2}
def test_automatic_create_constant_task(self):
with Flow(name="test") as flow:
t = Task()
t.set_dependencies(upstream_tasks=[4])
assert len(flow.tasks) == 2
assert any(isinstance(t, Constant) for t in flow.tasks)
def test_set_returns_a_set(self):
l = collections.Set()
with Flow(name="test") as f:
l.bind(1, 2)
assert f.run().result[l].result == set([1, 2])
def test_dict_returns_a_dict(self):
l = collections.Dict()
with Flow(name="test") as f:
l.bind(keys=["a", "b"], values=[1, 2])
assert f.run().result[l].result == dict(a=1, b=2)
def test_tuple_returns_a_tuple(self):
l = collections.Tuple()
with Flow(name="test") as f:
l.bind(1, 2)
assert f.run().result[l].result == (1, 2)
def test_initialize_returns_state_if_provided(self, state):
result = FlowRunner(Flow(name="test")).initialize_run(
state=state, task_states={}, context={}, task_contexts={}, parameters={}
)
assert result.state is state
def test_initialize_sets_none_to_pending(self):
result = FlowRunner(Flow(name="test")).initialize_run(
state=None, task_states={}, context={}, task_contexts={}, parameters={}
)
assert result.state.is_pending()
def test_scheduled_states_with_past_start_time(self):
state = Scheduled(start_time=pendulum.now("utc") -
datetime.timedelta(minutes=1))
assert (FlowRunner(flow=Flow(
name="test")).check_flow_reached_start_time(state=state) is state)
def test_tuple_maintains_sort_order_for_more_than_10_items(self):
# https://github.com/PrefectHQ/prefect/issues/2451
t = collections.Tuple()
with Flow(name="test") as f:
t.bind(*list(range(15)))
assert f.run().result[t].result == tuple(range(15))
def test_pending_becomes_running(self, state):
flow = Flow(name="test", tasks=[Task()])
new_state = FlowRunner(flow=flow).set_flow_to_running(state=state)
assert new_state.is_running()
def test_list_returns_a_list(self):
l = collections.List()
with Flow(name="test") as f:
l.bind(1, 2)
assert f.run().result[l].result == [1, 2]
def test_other_states_raise_endrun(self, state):
flow = Flow(name="test", tasks=[Task()])
with pytest.raises(ENDRUN):
FlowRunner(flow=flow).set_flow_to_running(state=state)
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",
}
def test_multiple_flow_handlers_are_called(self):
flow = Flow(name="test", state_handlers=[flow_handler, flow_handler])
FlowRunner(flow=flow).run()
# each flow changed state twice: Pending -> Running -> Success
assert handler_results["Flow"] == 4
def test_pending_or_running_are_ok(self, state):
flow = Flow(name="test", tasks=[Task()])
new_state = FlowRunner(flow=flow).check_flow_is_pending_or_running(
state=state)
assert new_state is state
def test_flow_runner_has_logger():
r = FlowRunner(Flow(name="test"))
assert r.logger.name == "prefect.FlowRunner"
def test_not_pending_or_running_raise_endrun(self, state):
flow = Flow(name="test", tasks=[Task()])
with pytest.raises(ENDRUN):
FlowRunner(flow=flow).check_flow_is_pending_or_running(state=state)
def test_outputs_task_decorator(self):
with Flow("test"):
assert self.mult(x=1).outputs() == int
def test_non_scheduled_states(self, state):
assert (FlowRunner(flow=Flow(
name="test")).check_flow_reached_start_time(state=state) is state)
def test_set_upstream_context(self):
with Flow(name="test") as f:
t1 = Task()
t2 = Task()
t2.set_upstream(t1)
assert Edge(t1, t2) in f.edges
def test_scheduled_states_without_start_time(self):
state = Scheduled(start_time=None)
assert (FlowRunner(flow=Flow(
name="test")).check_flow_reached_start_time(state=state) is state)
def test_set_upstream_with_properties(self, props):
with Flow(name="test") as f:
t1 = Task()
t2 = Task()
t2.set_upstream(t1, **props)
assert Edge(t1, t2, **props) in f.edges
def test_right_addition(self):
with Flow(name="test") as f:
z = 10 + Parameter("x")
state = f.run(parameters=dict(x=1))
assert state.result[z].result == 11