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()