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