def test_mapped_ifelse_and_merge():
@task
def is_even(x):
return x % 2 == 0
@task
def even():
return "even"
@task
def odd():
return "odd"
with Flow("iterated map") as flow:
mapped_result = is_even.map([1, 2, 3])
ifelse(condition=mapped_result,
true_task=even,
false_task=odd,
mapped=True)
merge_result = merge(even, odd, flow=flow, mapped=True)
state = flow.run()
assert state.result[merge_result].result == ["odd", "even", "odd"]
def test_list_of_tasks():
"""
Test that a list of tasks can be set as a switch condition
"""
with Flow(name="test") as flow:
condition = Condition()
true_branch = [SuccessTask(), SuccessTask()]
false_branch = SuccessTask()
ifelse(condition, true_branch, false_branch)
with prefect.context(CONDITION=True):
state = flow.run()
for t in true_branch:
assert isinstance(state.result[t], Success)
assert isinstance(state.result[false_branch], Skipped)
with prefect.context(CONDITION=False):
state = flow.run()
for t in true_branch:
# the tasks in the list ran becuase they have no upstream dependencies.
assert isinstance(state.result[t], Success)
list_task = next(t for t in flow.tasks
if isinstance(t, prefect.tasks.core.collections.List))
# but the list itself skipped
assert isinstance(state.result[list_task], Skipped)
assert isinstance(state.result[false_branch], Success)
def test_merging_diamond_flow():
"""
Test a flow that branches into two separate chains that later merge back together.
One branch should all get skipped but the merge task should not skip.
"""
with Flow(name="test") as flow:
condition = Condition()
true_branch = [
SuccessTask(name="true branch {}".format(i)) for i in range(3)
]
false_branch = [
SuccessTask(name="false branch {}".format(i)) for i in range(3)
]
ifelse(condition, true_branch[0], false_branch[0])
flow.chain(*true_branch)
flow.chain(*false_branch)
merge_task = merge(true_branch[-1], false_branch[-1])
with prefect.context(CONDITION=True):
state = flow.run()
for t in true_branch:
assert isinstance(state.result[t], Success)
for t in false_branch:
assert isinstance(state.result[t], Skipped)
assert isinstance(state.result[merge_task], Success)
def test_merge_diamond_flow_with_results():
condition = Condition()
@task
def true_branch():
return 1
@task
def false_branch():
return 0
with Flow(name="test") as flow:
ifelse(condition, true_branch, false_branch)
merge_task = merge(true_branch, false_branch)
with prefect.context(CONDITION=True):
state = flow.run()
assert state.result[merge_task].result == 1
with prefect.context(CONDITION=False):
state = flow.run()
assert state.result[merge_task].result == 0
with prefect.context(CONDITION=None):
state = flow.run()
assert state.result[merge_task].result is None
def main():
with Flow("manual-live-purpleair-control-flow") as flow:
offline = Parameter("offline", default=True)
all_sensors_online = extract_online_live_purpleair()
all_sensors_offline = extract_offline_live_purpleair()
ifelse(offline, all_sensors_offline, all_sensors_online)
all_sensors = merge(all_sensors_offline, all_sensors_online)
# Registers flow to server, which we can then deploy and run in background agents.
flow.register(project_name="caqi-flows")
def test_merge_can_distinguish_between_a_none_result_and_an_unrun_task():
condition = Condition()
true_branch = Constant(None)
false_branch = Constant(0)
with Flow(name="test") as flow:
ifelse(condition, true_branch, false_branch)
merge_task = merge(true_branch, false_branch)
with prefect.context(CONDITION=True):
state = flow.run()
assert state.result[merge_task].result is None
def test_merge_with_list():
with Flow(name="test") as flow:
condition = Condition()
true_branch = prefect.utilities.tasks.as_task([Constant(1), Constant(2)])
false_branch = Constant(0)
with pytest.warns(prefect.utilities.exceptions.PrefectWarning):
ifelse(condition, true_branch, false_branch)
merge_task = merge(true_branch, false_branch)
with prefect.context(CONDITION=True):
state = flow.run()
assert state.result[merge_task].result == [1, 2]
def test_ifelse_with_falsey_conditions(condition_value):
condition = Condition()
true_branch = SuccessTask(name="true branch")
false_branch = SuccessTask(name="false branch")
with Flow(name="test") as flow:
ifelse(condition, true_branch, false_branch)
with prefect.context(CONDITION=condition_value):
state = flow.run()
assert state.result[true_branch].is_skipped()
assert state.result[false_branch].is_successful()
def test_merge_with_list():
@task
def false_branch():
return 0
@task
def true_branch():
return [1, 2]
with Flow(name="test") as flow:
condition = Condition()
ifelse(condition, true_branch, false_branch)
merge_task = merge(true_branch, false_branch)
with prefect.context(CONDITION=True):
state = flow.run()
assert state.result[merge_task].result == [1, 2]
def test_merging_with_objects_that_cant_be_equality_compared():
class SpecialObject:
def __eq__(self, other):
return self
def __bool__(self):
raise SyntaxError("You can't handle the truth!")
@task
def return_array():
return SpecialObject()
with Flow("test-merge") as flow:
success = SuccessTask()
ifelse(Condition(), success, return_array)
merge_task = merge(success, return_array)
with prefect.context(CONDITION=False):
flow_state = flow.run()
assert flow_state.is_successful()
assert isinstance(flow_state.result[merge_task].result, SpecialObject)
def test_ifelse_doesnt_add_None_task(condition_value):
condition = Condition()
true_branch = SuccessTask(name="true branch")
with Flow(name="test") as flow:
cnd = ifelse(condition, true_branch, None)
assert len(flow.tasks) == 4
with prefect.context(CONDITION=condition_value):
state = flow.run()
assert isinstance(state.result[true_branch],
Success if condition_value is True else Skipped)
def test_ifelse(condition_value):
condition = Condition()
with Flow(name="test") as flow:
true_branch = identity("true")
false_branch = identity("false")
res = ifelse(condition, true_branch, false_branch)
assert len(flow.tasks) == 7
with prefect.context(CONDITION=condition_value):
state = flow.run()
if condition_value:
assert state.result[true_branch].is_successful()
assert state.result[false_branch].is_skipped()
assert state.result[res].result == "true"
else:
assert state.result[true_branch].is_skipped()
assert state.result[false_branch].is_successful()
assert state.result[res].result == "false"
def check_if_even(value):
return value % 2 == 0
@task
def print_odd(value):
print("{} is odd!".format(value))
@task
def print_even(value):
print("{} is even!".format(value))
with Flow("Check Even/Odd") as f:
value = Parameter("value")
is_even = check_if_even(value)
even = print_even(value)
odd = print_odd(value)
ifelse(is_even, even, odd)
# Prints '2 is even!'
f.run(value=2)
# Prints '1 is odd!'
f.run(value=1)
import prefect
from prefect import Flow, Task, task
from prefect.engine.result import NoResult
from prefect.engine.state import Skipped, Success
from prefect.tasks.control_flow import FilterTask, ifelse, merge, switch
from prefect.utilities.tasks import as_task
class Condition(Task):
def run(self):
return prefect.context.CONDITION
@task
def identity(x):
return x
condition = Condition()
with Flow(name="test") as flow:
true_branch = identity("true")
false_branch = identity("false")
res = ifelse(condition, true_branch, false_branch)
assert len(flow.tasks) == 7
with prefect.context(CONDITION=False):
state = flow.run()
