def test_task_runner_handles_version_lock_error(monkeypatch):
client = MagicMock()
monkeypatch.setattr("prefect.engine.cloud.task_runner.Client",
MagicMock(return_value=client))
client.set_task_run_state.side_effect = VersionLockError()
task = Task(name="test")
runner = CloudTaskRunner(task=task)
# successful state
client.get_task_run_state.return_value = Success()
res = runner.call_runner_target_handlers(Pending(), Running())
assert res.is_successful()
# currently running
client.get_task_run_state.return_value = Running()
with pytest.raises(ENDRUN):
runner.call_runner_target_handlers(Pending(), Running())
# result load error
s = Success()
s.load_result = MagicMock(side_effect=Exception())
client.get_task_run_state.return_value = s
with pytest.raises(ENDRUN):
res = runner.call_runner_target_handlers(Pending(), Running())
def vclient(monkeypatch):
cloud_client = MagicMock(
get_flow_run_info=MagicMock(return_value=MagicMock(state=None)),
set_flow_run_state=MagicMock(),
get_task_run_info=MagicMock(return_value=MagicMock(state=None)),
set_task_run_state=MagicMock(
side_effect=VersionLockError(),
return_value=Running()
# side_effect=lambda task_run_id, version, state, cache_for: state
),
get_latest_task_run_states=MagicMock(
side_effect=lambda flow_run_id, states: states),
)
monkeypatch.setattr("prefect.engine.cloud.task_runner.Client",
MagicMock(return_value=cloud_client))
monkeypatch.setattr("prefect.engine.cloud.flow_runner.Client",
MagicMock(return_value=cloud_client))
yield cloud_client
def test_flowrunner_handles_version_lock_error(monkeypatch):
client = MagicMock()
monkeypatch.setattr("prefect.engine.cloud.flow_runner.Client",
MagicMock(return_value=client))
client.set_flow_run_state.side_effect = VersionLockError()
flow = prefect.Flow(name="test")
runner = CloudFlowRunner(flow=flow)
# successful state
client.get_flow_run_state.return_value = Success()
res = runner.call_runner_target_handlers(Pending(), Running())
assert res.is_successful()
# currently running
client.get_flow_run_state.return_value = Running()
with pytest.raises(ENDRUN):
runner.call_runner_target_handlers(Pending(), Running())