def test_simple_map(monkeypatch, executor):
flow_run_id = str(uuid.uuid4())
task_run_id_1 = str(uuid.uuid4())
with prefect.Flow(name="test") as flow:
t1 = plus_one.map([0, 1, 2])
client = MockedCloudClient(
flow_runs=[FlowRun(id=flow_run_id)],
task_runs=[
TaskRun(id=task_run_id_1, task_id=t1.id, flow_run_id=flow_run_id)
] + [
TaskRun(id=t.id, task_id=t.id, flow_run_id=flow_run_id)
for t in flow.tasks if t is not t1
],
monkeypatch=monkeypatch,
)
with prefect.context(flow_run_id=flow_run_id):
state = CloudFlowRunner(flow=flow).run(return_tasks=flow.tasks,
executor=executor)
assert state.is_successful()
assert client.flow_runs[flow_run_id].state.is_successful()
assert client.task_runs[task_run_id_1].state.is_mapped()
# there should be a total of 4 task runs corresponding to the mapped task
assert len([tr for tr in client.task_runs.values()
if tr.task_id == t1.id]) == 4
def test_starting_at_arbitrary_loop_index_from_cloud_context(client):
@prefect.task
def looper(x):
if prefect.context.get("task_loop_count", 1) < 20:
raise LOOP(result=prefect.context.get("task_loop_result", 0) + x)
return prefect.context.get("task_loop_result", 0) + x
@prefect.task
def downstream(l):
return l ** 2
with prefect.Flow(name="looping") as f:
inter = looper(10)
final = downstream(inter)
client.get_flow_run_info = MagicMock(
return_value=MagicMock(context={"task_loop_count": 20})
)
client.set_flow_run_state = MagicMock()
flow_state = CloudFlowRunner(flow=f).run(return_tasks=[inter, final])
assert flow_state.is_successful()
assert flow_state.result[inter].result == 10
assert flow_state.result[final].result == 100
def test_task_failure_with_upstream_secrets_doesnt_store_secret_value_and_recompute_if_necessary(
client, ):
@prefect.task(max_retries=2, retry_delay=timedelta(minutes=100))
def is_p_three(p):
if p == 3:
raise ValueError("No thank you.")
return p
with prefect.Flow("test", result=PrefectResult()) as f:
p = prefect.tasks.secrets.PrefectSecret("p")
res = is_p_three(p)
with prefect.context(secrets=dict(p=3)):
state = CloudFlowRunner(flow=f).run(return_tasks=[res])
assert state.is_running()
assert isinstance(state.result[res], Retrying)
assert state.result[res].cached_inputs["p"].location is None
## here we set the result of the secret to an empty result, ensuring
## it will get converted to a "true" result;
## we expect that the upstream value will actually get recomputed from context
## through the SecretResultHandler
safe = SecretResult(p)
state.result[p] = Success(result=safe)
state.result[res].start_time = pendulum.now("utc")
state.result[res].cached_inputs = dict(p=safe)
with prefect.context(secrets=dict(p=4)):
new_state = CloudFlowRunner(flow=f).run(return_tasks=[res],
task_states=state.result)
assert new_state.is_successful()
assert new_state.result[res].result == 4
def test_simple_two_task_flow(monkeypatch, executor):
flow_run_id = str(uuid.uuid4())
task_run_id_1 = str(uuid.uuid4())
task_run_id_2 = str(uuid.uuid4())
with prefect.Flow(name="test") as flow:
t1 = prefect.Task()
t2 = prefect.Task()
t2.set_upstream(t1)
client = MockedCloudClient(
flow_runs=[FlowRun(id=flow_run_id)],
task_runs=[
TaskRun(id=task_run_id_1, task_id=t1.id, flow_run_id=flow_run_id),
TaskRun(id=task_run_id_2, task_id=t2.id, flow_run_id=flow_run_id),
],
monkeypatch=monkeypatch,
)
with prefect.context(flow_run_id=flow_run_id):
state = CloudFlowRunner(flow=flow).run(return_tasks=flow.tasks,
executor=executor)
assert state.is_successful()
assert client.flow_runs[flow_run_id].state.is_successful()
assert client.task_runs[task_run_id_1].state.is_successful()
assert client.task_runs[task_run_id_1].version == 2
assert client.task_runs[task_run_id_2].state.is_successful()
def test_scheduled_start_time_is_in_context(monkeypatch, executor):
flow_run_id = str(uuid.uuid4())
task_run_id_1 = str(uuid.uuid4())
flow = prefect.Flow(name="test", tasks=[whats_the_time])
client = MockedCloudClient(
flow_runs=[FlowRun(id=flow_run_id)],
task_runs=[
TaskRun(
id=task_run_id_1, task_slug=whats_the_time.slug, flow_run_id=flow_run_id
)
],
monkeypatch=monkeypatch,
)
with prefect.context(flow_run_id=flow_run_id):
state = CloudFlowRunner(flow=flow).run(
return_tasks=flow.tasks, executor=executor
)
assert state.is_successful()
assert client.flow_runs[flow_run_id].state.is_successful()
assert client.task_runs[task_run_id_1].state.is_successful()
assert isinstance(state.result[whats_the_time].result, datetime.datetime)
def test_deep_map(monkeypatch, executor):
flow_run_id = str(uuid.uuid4())
task_run_id_1 = str(uuid.uuid4())
task_run_id_2 = str(uuid.uuid4())
task_run_id_3 = str(uuid.uuid4())
with prefect.Flow(name="test", result_handler=JSONResultHandler()) as flow:
t1 = plus_one.map([0, 1, 2])
t2 = plus_one.map(t1)
t3 = plus_one.map(t2)
client = MockedCloudClient(
flow_runs=[FlowRun(id=flow_run_id)],
task_runs=[
TaskRun(
id=task_run_id_1, task_slug=flow.slugs[t1], flow_run_id=flow_run_id
),
TaskRun(
id=task_run_id_2, task_slug=flow.slugs[t2], flow_run_id=flow_run_id
),
TaskRun(
id=task_run_id_3, task_slug=flow.slugs[t3], flow_run_id=flow_run_id
),
]
+ [
TaskRun(
id=str(uuid.uuid4()), task_slug=flow.slugs[t], flow_run_id=flow_run_id
)
for t in flow.tasks
if t not in [t1, t2, t3]
],
monkeypatch=monkeypatch,
)
with prefect.context(flow_run_id=flow_run_id):
state = CloudFlowRunner(flow=flow).run(
return_tasks=flow.tasks, executor=executor
)
assert state.is_successful()
assert client.flow_runs[flow_run_id].state.is_successful()
assert client.task_runs[task_run_id_1].state.is_mapped()
assert client.task_runs[task_run_id_2].state.is_mapped()
assert client.task_runs[task_run_id_3].state.is_mapped()
# there should be a total of 4 task runs corresponding to each mapped task
for t in [t1, t2, t3]:
assert (
len(
[
tr
for tr in client.task_runs.values()
if tr.task_slug == flow.slugs[t]
]
)
== 4
)
def test_flow_runner_puts_running_with_backend_in_context(client):
@prefect.task()
def whats_in_ctx():
assert prefect.context.get("running_with_backend")
flow = prefect.Flow(name="test", tasks=[whats_in_ctx])
res = CloudFlowRunner(flow=flow).run()
assert res.is_successful()
def test_can_queue_successfully_and_run(monkeypatch):
@prefect.task
def return_one():
return 1
with prefect.Flow("test-queues-work!") as flow:
t1 = return_one()
flow_run_id = str(uuid.uuid4())
task_run_id_1 = str(uuid.uuid4())
client = QueueingMockCloudClient(
flow_runs=[FlowRun(id=flow_run_id)],
task_runs=[
TaskRun(
id=task_run_id_1, task_slug=flow.slugs[t1], flow_run_id=flow_run_id
),
]
+ [
TaskRun(
id=str(uuid.uuid4()), task_slug=flow.slugs[t1], flow_run_id=flow_run_id
)
for t in flow.tasks
if t
not in [
t1,
]
],
monkeypatch=monkeypatch,
num_times_in_queue=6,
)
with prefect.context(flow_run_id=flow_run_id):
run_state = CloudFlowRunner(flow=flow).run(
executor=LocalExecutor(), return_tasks=flow.tasks
)
assert run_state.is_successful()
# Pending -> Running -> Queued (4x) -> Success
# State transitions that result in `set_flow_run_state` calls are from
# Pending -> Running and Running -> Success, all others
# are from Running -> Queued or Queued -> Queued
assert client.call_count["set_flow_run_state"] == 2 + (client.num_times_in_queue)
def test_flow_runner_retries_forever_on_queued_state(client, monkeypatch,
num_attempts):
mock_sleep = MagicMock()
monkeypatch.setattr("prefect.engine.cloud.flow_runner.time.sleep",
mock_sleep)
run_states = [
Queued(start_time=pendulum.now("UTC").add(seconds=i))
for i in range(num_attempts - 1)
]
run_states.append(Success())
mock_run = MagicMock(side_effect=run_states)
client.get_flow_run_info = MagicMock(
side_effect=[MagicMock(version=i) for i in range(num_attempts)])
# Mock out the actual flow execution
monkeypatch.setattr("prefect.engine.cloud.flow_runner.FlowRunner.run",
mock_run)
@prefect.task
def return_one():
return 1
with prefect.Flow("test-cloud-flow-runner-with-queues") as flow:
one = return_one()
# Without these (actual, not mocked) sleep calls, when running full test suite this
# test can fail for no reason.
final_state = CloudFlowRunner(flow=flow).run()
assert final_state.is_successful()
assert mock_run.call_count == num_attempts
# Not called on the initial run attempt
assert client.get_flow_run_info.call_count == num_attempts - 1
