def test_all_pipeline_method_steps_are_called():
pipeline = [
"initialize_run",
"check_flow_is_pending_or_running",
"set_flow_to_running",
"get_flow_run_state",
]
runner = FlowRunner(Flow(name="test"))
for method in pipeline:
setattr(runner, method, MagicMock())
# initialize run is unpacked, which MagicMocks dont support
runner.initialize_run = MagicMock(
return_value=FlowRunnerInitializeResult(
MagicMock(), MagicMock(), MagicMock(), MagicMock()
)
)
runner.run()
for method in pipeline:
assert getattr(runner, method).call_count == 1
def test_constant_tasks_arent_submitted_when_mapped(caplog):
calls = []
class TrackSubmissions(LocalExecutor):
def submit(self, *args, **kwargs):
calls.append(kwargs)
return super().submit(*args, **kwargs)
@prefect.task
def add(x):
return x + 1
with Flow("constants") as flow:
output = add.map([99] * 10)
runner = FlowRunner(flow=flow)
flow_state = runner.run(return_tasks=[output], executor=TrackSubmissions())
assert flow_state.is_successful()
assert flow_state.result[output].result == [100] * 10
## the add task was submitted 11 times: one for the parent and 10 times for each child
assert len(calls) == 11
## to be safe, ensure '5' isn't in the logs
assert len([log.message for log in caplog.records
if "99" in log.message]) == 0
def test_constant_tasks_arent_submitted(caplog):
calls = []
class TrackSubmissions(LocalExecutor):
def submit(self, *args, **kwargs):
calls.append(kwargs)
return super().submit(*args, **kwargs)
@prefect.task
def add(x):
return x + 1
with Flow("constants") as flow:
output = add(5)
runner = FlowRunner(flow=flow)
flow_state = runner.run(return_tasks=[output], executor=TrackSubmissions())
assert flow_state.is_successful()
assert flow_state.result[output].result == 6
## only add was submitted
assert len(calls) == 1
## to be safe, ensure '5' isn't in the logs
assert len([log.message for log in caplog.records
if "5" in log.message]) == 0
def test_flow_runner_with_invalid_return_tasks():
flow = Flow(name="test")
task = SuccessTask()
flow.add_task(task)
flow_runner = FlowRunner(flow=flow)
state = flow_runner.run(return_tasks=[1])
assert state.is_failed()
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_overwrite_context():
x = prefect.Parameter("x")
f = FlowRunner(Flow(name="test", tasks=[x]))
state = f.run(
parameters={"x": 2}, context={"parameters": {"x": 5}}, return_tasks=[x]
)
assert state.result[x].result == 2
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_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_parameters_overwrite_context_only_if_key_matches():
x = prefect.Parameter("x")
y = prefect.Parameter("y")
f = FlowRunner(Flow(name="test", tasks=[x, y]))
state = f.run(
parameters={"x": 2},
context={"parameters": {"x": 5, "y": 6}},
return_tasks=[x, y],
)
def test_flow_runner_passes_along_its_run_context_to_tasks(self):
@prefect.task
def grab_key():
return prefect.context["THE_ANSWER"]
with prefect.context(THE_ANSWER=42):
runner = FlowRunner(Flow(name="test", tasks=[grab_key]))
flow_state = runner.run(return_tasks=[grab_key])
assert flow_state.is_successful()
assert flow_state.result[grab_key].result == 42
def test_map_preserves_flowrunners_run_context(executor):
@task
def whats_id():
return prefect.context.get("special_id")
with Flow(name="test-context-preservation") as flow:
result = whats_id.map(upstream_tasks=[list(range(10))])
with prefect.context(special_id="FOOBAR"):
runner = FlowRunner(flow=flow)
flow_state = runner.run(return_tasks=[result])
assert flow_state.is_successful()
assert flow_state.result[result].result == ["FOOBAR"] * 10
