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 = VersionLockMismatchSignal() 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 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 = VersionLockMismatchSignal() 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())
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=VersionLockMismatchSignal(), 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