def test_context_includes_date(self): @prefect.task def return_ctx_key(): return prefect.context.get("date") f = Flow(name="test", tasks=[return_ctx_key]) res = f.run() assert res.is_successful() output = res.result[return_ctx_key].result assert isinstance(output, datetime.datetime)
def test_user_provided_context_is_prioritized(self): @prefect.task def return_ctx_key(): return prefect.context.get("date") f = Flow(name="test", tasks=[return_ctx_key]) res = f.run(context={"date": "42"}) assert res.is_successful() output = res.result[return_ctx_key].result assert output == "42"
def test_flow_runner_makes_copy_of_task_results_dict(): """ Ensure the flow runner copies the task_results dict rather than modifying it inplace """ flow = Flow(name="test") t1, t2 = Task(), Task() flow.add_edge(t1, t2) task_states = {t1: Pending()} state = flow.run(task_states=task_states) assert state.result[t1] == Success(result=None) assert task_states == {t1: Pending()}
def test_flow_run_method_returns_task_states_even_if_it_doesnt_run(): # https://github.com/PrefectHQ/prefect/issues/19 flow = Flow(name="test") task1 = SuccessTask() task2 = ErrorTask() flow.add_edge(task1, task2) flow_state = flow.run(state=Success()) assert flow_state.is_successful() assert flow_state.result[task1].is_pending() assert flow_state.result[task2].is_pending()
def test_flow_runner_doesnt_override_scheduled_start_time_when_running_on_schedule( self, run_on_schedule): @prefect.task def return_scheduled_start_time(): return prefect.context.get("scheduled_start_time") f = Flow(name="test", tasks=[return_scheduled_start_time]) res = f.run(context=dict(scheduled_start_time=42), run_on_schedule=run_on_schedule) assert res.is_successful() assert res.result[return_scheduled_start_time].result != 42
def test_flow_runner_provides_scheduled_start_time(self): @prefect.task def return_scheduled_start_time(): return prefect.context.get("scheduled_start_time") f = Flow(name="test", tasks=[return_scheduled_start_time]) res = f.run() assert res.is_successful() assert res.result[return_scheduled_start_time].is_successful() assert isinstance(res.result[return_scheduled_start_time].result, datetime.datetime)
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_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_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_context_contains_date_formats(self, date): @prefect.task def return_ctx_key(): return prefect.context.get(date) f = Flow(name="test", tasks=[return_ctx_key]) res = f.run() assert res.is_successful() output = res.result[return_ctx_key].result assert isinstance(output, str) assert len(output) == 10
def test_flow_runner_captures_and_exposes_dask_errors(executor): q = queue.Queue() @prefect.task def put(): q.put(55) f = Flow(name="test", tasks=[put]) state = f.run(executor=executor) assert state.is_failed() assert isinstance(state.result, TypeError) assert str(state.result) == "can't pickle _thread.lock objects"
def test_user_provided_context_is_prioritized(self, outer_context, inner_context, sol): @prefect.task def return_ctx_key(): return prefect.context.get("date") f = Flow(name="test", tasks=[return_ctx_key]) with prefect.context(**outer_context): res = f.run(context=inner_context) assert res.is_successful() output = res.result[return_ctx_key].result assert output == sol
def test_flow_runner_captures_and_exposes_dask_errors(executor): q = queue.Queue() @prefect.task def put(): q.put(55) f = Flow(name="test", tasks=[put]) state = f.run(executor=executor) assert state.is_failed() assert isinstance(state.result, TypeError) # assert two possible result outputs for different Python versions assert str(state.result) in [ "can't pickle _thread.lock objects", "cannot pickle '_thread.lock' object", ]
def test_flow_runner_properly_provides_context_to_task_runners(executor): @prefect.task def my_name(): return prefect.context.get("my_name") @prefect.task def flow_name(): return prefect.context.get("flow_name") flow = Flow(name="test-dummy", tasks=[flow_name, my_name]) with prefect.context(my_name="marvin"): res = flow.run(executor=executor) assert res.result[flow_name].result == "test-dummy" assert res.result[my_name].result == "marvin" with Flow("test-map") as f: tt = flow_name.map(upstream_tasks=[my_name]) with prefect.context(my_name="mapped-marvin"): res = f.run(executor=executor) assert res.result[my_name].result == "mapped-marvin" assert res.result[tt].result[0] == "test-map"
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()
def test_task_runners_submitted_to_remote_machines_respect_original_config( monkeypatch): """ This test is meant to simulate the behavior of running a Cloud Flow against an external cluster which has _not_ been configured for Prefect. The idea is that the configuration settings which were present on the original machine are respected in the remote job, reflected here by having the CloudHandler called during logging and the special values present in context. """ from prefect.engine.flow_runner import run_task def my_run_task(*args, **kwargs): with prefect.utilities.configuration.set_temporary_config({ "logging.log_to_cloud": False, "cloud.auth_token": "" }): return run_task(*args, **kwargs) calls = [] class Client: def write_run_logs(self, *args, **kwargs): calls.append(args) monkeypatch.setattr("prefect.engine.flow_runner.run_task", my_run_task) monkeypatch.setattr("prefect.client.Client", Client) @prefect.task def log_stuff(): logger = prefect.context.get("logger") logger.critical("important log right here") return ( prefect.context.config.special_key, prefect.context.config.cloud.auth_token, ) with prefect.utilities.configuration.set_temporary_config({ "logging.log_to_cloud": True, "special_key": 42, "cloud.auth_token": "original", }): # captures config at init flow = Flow("test", tasks=[log_stuff]) flow_state = flow.run(task_contexts={log_stuff: dict(special_key=99)}) assert flow_state.is_successful() assert flow_state.result[log_stuff].result == (42, "original") time.sleep(0.75) assert len(calls) >= 1 assert len([log for call in calls for log in call[0]]) == 5 # actual number of logs loggers = [log["name"] for call in calls for log in call[0]] assert set(loggers) == { "prefect.TaskRunner", "prefect.FlowRunner", "prefect.log_stuff", }