def test_map_allows_retries_2(executor):
"""
Another test of mapping and retries
"""
@prefect.task
def ll():
return [0, 1, 2]
div = DivTask(max_retries=1, retry_delay=datetime.timedelta(0))
with Flow(name="test") as f:
res = div.map(x=ll)
s = FlowRunner(flow=f).run(executor=executor, return_tasks=f.tasks)
assert s.is_running()
m = s.result[res]
assert m.map_states[0].is_pending()
assert m.map_states[1].is_successful()
assert m.map_states[2].is_successful()
s.result[ll].result[0] = 10
s = FlowRunner(flow=f).run(executor=executor,
task_states=s.result,
return_tasks=f.tasks)
assert s.is_successful()
assert s.result[res].result[0] == 1 / 10
def test_task_map_with_no_upstream_results_and_a_mapped_state(executor):
"""
This test makes sure that mapped tasks properly generate children tasks even when
run multiple times and without available upstream results. In this test, we run the pipeline
from a variety of starting points, ensuring that some upstream results are unavailable and
checking that children pipelines are properly regenerated.
Note that upstream results will be hydrated from remote locations when running with a Cloud TaskRunner.
"""
@prefect.task
def numbers():
return [1, 2, 3]
@prefect.task
def identity(x):
return x
with Flow(name="test") as f:
n = numbers()
x = identity.map(n)
# first run with a missing result from `n` but map_states for `x`
state = FlowRunner(flow=f).run(
executor=executor,
task_states={
n: Success(),
x: Mapped(map_states=[Pending() for i in range(1, 4)]),
},
return_tasks=f.tasks,
)
assert state.is_successful()
assert state.result[x].result == [None] * 3
def test_map_allows_for_retries(executor):
ii = IdTask()
ll = ListTask()
div = DivTask(max_retries=1, retry_delay=datetime.timedelta(0))
with Flow(name="test") as f:
l_res = ll(start=0)
divved = div.map(l_res)
res = ii.map(divved)
states = FlowRunner(flow=f).run(executor=executor, return_tasks=f.tasks)
assert states.is_running() # division by zero caused map to retry
old = states.result[divved]
assert old.result[1:] == [1.0, 0.5]
assert old.map_states[0].is_retrying()
# update upstream result
states.result[l_res].result[0] = 0.01
states = FlowRunner(flow=f).run(task_states=states.result,
executor=executor,
return_tasks=f.tasks)
assert states.is_successful() # no divison by 0
new = states.result[res]
assert new.result == [100, 1.0, 0.5]
def test_flow_handlers_can_return_none(self):
flow_handler = MagicMock(side_effect=lambda t, o, n: None)
flow = Flow(name="test", state_handlers=[flow_handler])
flow_state = FlowRunner(flow=flow).run()
assert flow_state.is_successful()
# the flow changed state twice: Pending -> Running -> Success
assert flow_handler.call_count == 2
def test_running_state_finishes(self):
flow = Flow(name="test", tasks=[Task()])
new_state = FlowRunner(flow=flow).get_flow_run_state(
state=Running(),
task_states={},
task_contexts={},
return_tasks=set(),
task_runner_state_handlers=[],
executor=LocalExecutor(),
)
assert new_state.is_successful()
def test_mapped_task_can_be_scheduled(self, executor):
with Flow(name="test") as flow:
res = ReturnTask().map([0, 0])
state = FlowRunner(flow=flow).run(
return_tasks=[res],
executor=executor,
task_states={res: Scheduled(start_time=pendulum.now().subtract(minutes=1))},
)
assert state.is_successful()
def test_mapped_will_use_existing_map_states_if_available(self, executor):
with Flow(name="test") as flow:
res = ReturnTask().map([0, 1])
state = FlowRunner(flow=flow).run(
return_tasks=[res],
executor=executor,
task_states={res: Mapped(map_states=[Success(), Success(result=100)])},
)
assert state.is_successful()
assert state.result[res].map_states[1].is_successful()
assert state.result[res].map_states[1].result == 100
def test_secrets_are_rerun_on_restart():
@prefect.task
def identity(x):
return x
with Flow("test") as flow:
secret = PrefectSecret("key")
val = identity(secret)
with prefect.context(secrets={"key": "val"}):
state = FlowRunner(flow=flow).run(task_states={secret: Success()},
return_tasks=[val])
assert state.is_successful()
assert state.result[val].result == "val"
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_as_task_with_basic_python_objs(self, obj):
@tasks.task
def return_val(x):
"Necessary because constant tasks aren't tracked inside the flow"
return x
with Flow("test") as f:
t = tasks.as_task(obj)
val = return_val(t)
assert isinstance(t, Task)
res = FlowRunner(f).run(return_tasks=[val])
assert res.is_successful()
assert res.result[val].result == obj
def test_task_map_with_no_upstream_results_and_a_mapped_state(executor):
"""
This test makes sure that mapped tasks properly generate children tasks even when
run multiple times and without available upstream results. In this test, we run the pipeline
from a variety of starting points, ensuring that some upstream results are unavailable and
checking that children pipelines are properly regenerated.
"""
@prefect.task
def numbers():
return [1, 2, 3]
@prefect.task
def plus_one(x):
return x + 1
@prefect.task
def get_sum(x):
return sum(x)
with Flow(name="test") as f:
n = numbers()
x = plus_one.map(n)
y = plus_one.map(x)
s = get_sum(y)
# first run with a missing result from `n` but map_states for `x`
state = FlowRunner(flow=f).run(
executor=executor,
task_states={
n:
Success(),
x:
Mapped(map_states=[
Pending(cached_inputs={"x": Result(i)}) for i in range(1, 4)
]),
},
return_tasks=f.tasks,
)
assert state.is_successful()
assert state.result[s].result == 12
# next run with missing results for n and x
state = FlowRunner(flow=f).run(
executor=executor,
task_states={
n:
Success(),
x:
Mapped(map_states=[Success(), Success(),
Success()]),
y:
Mapped(map_states=[
Success(result=3),
Success(result=4),
Retrying(cached_inputs={"x": Result(4)}),
]),
},
return_tasks=f.tasks,
)
assert state.is_successful()
assert state.result[s].result == 12
# next run with missing results for n, x, and y
state = FlowRunner(flow=f).run(
executor=executor,
task_states={
n:
Success(),
x:
Mapped(map_states=[Success(), Success(),
Success()]),
y:
Mapped(map_states=[
Success(result=3),
Success(result=4),
Success(result=5)
]),
},
return_tasks=f.tasks,
)
assert state.is_successful()
assert state.result[s].result == 12
