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_are_placed_into_context(): flow = Flow(name="test") y = prefect.Parameter("y", default=99) flow.add_task(y) flow_state = FlowRunner(flow=flow).run(return_tasks=[y], parameters=dict(y=42)) assert isinstance(flow_state, Success) assert flow_state.result[y].result == 42
def test_required_parameters_must_be_provided(): flow = Flow(name="test") y = prefect.Parameter("y") flow.add_task(y) flow_state = FlowRunner(flow=flow).run(return_tasks=[y]) assert isinstance(flow_state, Failed) assert isinstance(flow_state.result[y], Failed) assert "required but not provided" in str(flow_state.result[y]).lower()
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_deserialize_with_parameters_key(): f = Flow(name="test") x = Parameter("x") f.add_task(x) f2 = FlowSchema().load(FlowSchema().dump(f)) assert {p.name for p in f2.parameters()} == {p.name for p in f.parameters()} f_params = {(p.name, p.required, p.default) for p in f.parameters()} f2_params = {(p.name, p.required, p.default) for p in f2.parameters()} assert f_params == f2_params
def test_raise_on_exception_plays_well_with_context(): flow = Flow(name="test") flow.add_task(MathTask()) try: assert "raise_on_exception" not in prefect.context with raise_on_exception(): assert "raise_on_exception" in prefect.context flow.run() except ZeroDivisionError: assert "raise_on_exception" not in prefect.context pass
def test_nested_collection_automatically_applied_to_callargs_imperative(self): x = Parameter("x") y = Parameter("y") identity = IdentityTask() f = Flow(name="test") f.add_task(identity) identity.bind(x=dict(a=[x, dict(y=y)], b=(y, set([x]))), flow=f) state = f.run(parameters=dict(x=1, y=2)) assert len(f.tasks) == 10 assert state.result[identity].result == dict(a=[1, dict(y=2)], b=(2, set([1])))
def test_set_automatically_applied_to_callargs_imperative(self): x = Parameter("x") y = Parameter("y") identity = IdentityTask() f = Flow(name="test") f.add_task(identity) identity.bind(x=set([x, y]), flow=f) state = f.run(parameters=dict(x=1, y=2)) assert len(f.tasks) == 4 assert sum(isinstance(t, collections.Set) for t in f.tasks) == 1 assert state.result[identity].result == set([1, 2])
def test_dict_automatically_applied_to_callargs_imperative(self): x = Parameter("x") y = Parameter("y") identity = IdentityTask() f = Flow(name="test") f.add_task(identity) identity.bind(x=dict(a=x, b=y), flow=f) state = f.run(parameters=dict(x=1, y=2)) assert len(f.tasks) == 5 # 2 params, identity, Dict, List of dict values assert sum(isinstance(t, collections.Dict) for t in f.tasks) == 1 assert state.result[identity].result == dict(a=1, b=2)
def test_flow_runner_runs_basic_flow_with_2_independent_tasks(): flow = Flow(name="test") task1 = SuccessTask() task2 = SuccessTask() flow.add_task(task1) flow.add_task(task2) flow_state = FlowRunner(flow=flow).run(return_tasks=[task1, task2]) assert isinstance(flow_state, Success) assert flow_state.result[task1] == Success(result=1) assert flow_state.result[task2] == Success(result=1)
def test_run_count_updates_after_each_retry(self): flow = Flow(name="test") t1 = ErrorTask(max_retries=2, retry_delay=datetime.timedelta(0)) flow.add_task(t1) state1 = FlowRunner(flow=flow).run(return_tasks=[t1]) assert state1.result[t1].is_retrying() assert state1.result[t1].run_count == 1 state2 = FlowRunner(flow=flow).run(return_tasks=[t1], task_states=state1.result) assert state2.result[t1].is_retrying() assert state2.result[t1].run_count == 2
def test_parameters(): f = Flow(name="test") x = Parameter("x") y = Parameter("y", default=5) f.add_task(x) f.add_task(y) serialized = FlowSchema().dump(f) assert "parameters" in serialized assert [ isinstance(ParameterSchema().load(p), Parameter) for p in serialized["parameters"] ]
def test_raise_on_exception_ignores_all_prefect_signals(signal): flow = Flow(name="test") @prefect.task def raise_signal(): if (prefect.context.get("task_loop_count", 1) < 2 and prefect.context.get("task_run_count", 1) < 2 and signal.__name__ != "PAUSE"): raise signal("my message") flow.add_task(raise_signal) with raise_on_exception(): flow_state = flow.run()
def test_that_bad_code_in_flow_runner_is_caught(): """ Test that an error in the actual FlowRunner code itself (not the execution of user code) is caught """ class BadFlowRunner(FlowRunner): def get_flow_run_state(self, *args, **kwargs): raise RuntimeError("I represent bad code in the flow runner.") flow = Flow(name="test") flow.add_task(MathTask()) flow_state = BadFlowRunner(flow=flow).run() assert isinstance(flow_state, state.Failed) assert isinstance(flow_state.result, RuntimeError) assert "I represent bad code in the flow runner." == str(flow_state.result)
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_that_bad_code_in_task_runner_is_caught(): """ Test that an error in the actual TaskRunner code itself (not the execution of user code) is caught """ class BadTaskRunner(TaskRunner): def get_task_run_state(self, *args, **kwargs): raise RuntimeError("I represent bad code in the task runner.") flow = Flow(name="test") math_task = MathTask() flow.add_task(math_task) flow_state = FlowRunner( flow=flow, task_runner_cls=BadTaskRunner).run(return_tasks=[math_task]) assert isinstance(flow_state, state.Failed) assert flow_state.result[math_task].is_failed() caught_exc = flow_state.result[math_task].result assert isinstance(caught_exc, RuntimeError) assert "I represent bad code in the task runner." == str(caught_exc)
def test_raise_on_exception_raises_basic_error(): flow = Flow(name="test") flow.add_task(MathTask()) with pytest.raises(ZeroDivisionError): with raise_on_exception(): flow.run()