def test_scheduled_start_time_is_in_context(monkeypatch, executor): flow_run_id = str(uuid.uuid4()) task_run_id_1 = str(uuid.uuid4()) flow = prefect.Flow(name="test", tasks=[whats_the_time], result_handler=ResultHandler()) client = MockedCloudClient( flow_runs=[FlowRun(id=flow_run_id)], task_runs=[ TaskRun( id=task_run_id_1, task_slug=flow.slugs[whats_the_time], flow_run_id=flow_run_id, ) ], monkeypatch=monkeypatch, ) with prefect.context(flow_run_id=flow_run_id): state = CloudFlowRunner(flow=flow).run(return_tasks=flow.tasks, executor=executor) assert state.is_successful() assert client.flow_runs[flow_run_id].state.is_successful() assert client.task_runs[task_run_id_1].state.is_successful() assert isinstance(state.result[whats_the_time].result, datetime.datetime)
def test_task_runner_puts_cloud_in_context(client): @prefect.task(result_handler=ResultHandler()) def whats_in_ctx(): return prefect.context.get("checkpointing") res = CloudTaskRunner(task=whats_in_ctx).run() assert res.is_successful() assert res.result is True
def test_cloud_task_runner_handles_retries_with_queued_states_from_cloud(client): calls = [] def queued_mock(*args, **kwargs): calls.append(kwargs) # first retry attempt will get queued if len(calls) == 4: return Queued() # immediate start time else: return kwargs.get("state") client.set_task_run_state = queued_mock @prefect.task( max_retries=2, retry_delay=datetime.timedelta(seconds=0), result_handler=ResultHandler(), ) def tagged_task(x): if prefect.context.get("task_run_count", 1) == 1: raise ValueError("gimme a sec") return x upstream_result = Result(value=42, result_handler=JSONResultHandler()) res = CloudTaskRunner(task=tagged_task).run( context={"task_run_version": 1}, state=None, upstream_states={ Edge(Task(), tagged_task, key="x"): Success(result=upstream_result) }, executor=prefect.engine.executors.LocalExecutor(), ) assert res.is_successful() assert res.result == 42 assert ( len(calls) == 6 ) # Running -> Failed -> Retrying -> Queued -> Running -> Success assert [type(c["state"]).__name__ for c in calls] == [ "Running", "Failed", "Retrying", "Running", "Running", "Success", ] # ensures result handler was called and persisted assert calls[2]["state"].cached_inputs["x"].safe_value.value == "42"
def test_task_runner_handles_looping_with_no_result(client): @prefect.task(result_handler=ResultHandler()) def looper(): if prefect.context.get("task_loop_count", 1) < 3: raise LOOP() return 42 res = CloudTaskRunner(task=looper).run(context={"task_run_version": 1}, state=None, upstream_states={}) ## assertions assert res.is_successful() assert client.get_task_run_info.call_count == 0 assert ( client.set_task_run_state.call_count == 6 ) # Pending -> Running -> Looped (1) -> Running -> Looped (2) -> Running -> Success versions = [ call[1]["version"] for call in client.set_task_run_state.call_args_list if call[1]["version"] ] assert versions == [1, 3, 5]
def test_task_runner_handles_looping(client): @prefect.task(result_handler=ResultHandler()) def looper(): if prefect.context.get("task_loop_count", 1) < 3: raise LOOP(result=prefect.context.get("task_loop_result", 0) + 10) return prefect.context.get("task_loop_result") res = CloudTaskRunner(task=looper).run( context={"task_run_version": 1}, state=None, upstream_states={}, executor=prefect.engine.executors.LocalExecutor(), ) ## assertions assert res.is_successful() assert client.get_task_run_info.call_count == 0 assert ( client.set_task_run_state.call_count == 6 ) # Pending -> Running -> Looped (1) -> Running -> Looped (2) -> Running -> Success versions = [call[1]["version"] for call in client.set_task_run_state.call_args_list] assert versions == [1, 2, 3, 4, 5, 6]
def __init__(self) -> None: super().__init__(value=None, result_handler=ResultHandler())
def create_object(self, data: dict, **kwargs: Any) -> ResultHandler: """Because we cannot deserialize a custom class, just return `None`""" return ResultHandler()
def test_result_handler_base_class_is_a_passthrough(): handler = ResultHandler() assert handler.write("foo") == "foo" assert handler.read(99) == 99
def __init__(self) -> None: self.flows = dict() # type: Dict[str, prefect.core.flow.Flow] super().__init__(result_handler=ResultHandler())
def create_object(self, data: dict, **kwargs: Any) -> results.ResultHandlerResult: data["result_handler"] = ResultHandler() base_obj = super().create_object(data) return base_obj
def test_noresult_has_base_handler(): n = NoResult n.result_handler == ResultHandler()
def __init__(self) -> None: self.flows = dict() # type: Dict[str, bytes] super().__init__(result_handler=ResultHandler())
def test_purgedresult_has_base_handler(): n = PurgedResult n.result_handler == ResultHandler()