def test_heartbeat_traps_errors_caused_by_client(self, monkeypatch):
client = MagicMock(update_task_run_heartbeat=MagicMock(side_effect=SyntaxError))
monkeypatch.setattr(
"prefect.engine.cloud.task_runner.Client", MagicMock(return_value=client)
)
runner = CloudTaskRunner(task=Task(name="bad"))
runner.task_run_id = None
with pytest.warns(UserWarning) as warning:
res = runner._heartbeat()
assert res is None
assert client.update_task_run_heartbeat.called
w = warning.pop()
assert "Heartbeat failed for Task 'bad'" in repr(w.message)
def test_heartbeat_traps_errors_caused_by_client(self, caplog,
monkeypatch):
client = MagicMock(graphql=MagicMock(side_effect=SyntaxError))
monkeypatch.setattr("prefect.engine.cloud.task_runner.Client",
MagicMock(return_value=client))
runner = CloudTaskRunner(task=Task(name="bad"))
runner.task_run_id = None
res = runner._heartbeat()
assert res is False
log = caplog.records[0]
assert log.levelname == "ERROR"
assert "Heartbeat failed for Task 'bad'" in log.message
def multiprocessing_helper(executor):
set_task_run_state = MagicMock(
side_effect=lambda task_run_id, version, state, cache_for:
state)
client = MagicMock(set_task_run_state=set_task_run_state)
monkeypatch.setattr(
"prefect.engine.cloud.task_runner.Client",
MagicMock(return_value=client),
)
runner = CloudTaskRunner(task=Task())
runner.cache_result = lambda *args, **kwargs: time.sleep(0.25)
runner._heartbeat = update
with set_temporary_config({"cloud.heartbeat_interval": 0.05}):
return runner.run(executor=executor)
def test_task_runner_heartbeat_sets_command(self, monkeypatch):
client = MagicMock()
monkeypatch.setattr("prefect.engine.cloud.task_runner.Client",
MagicMock(return_value=client))
client.graphql.return_value.data.flow_run_by_pk.flow.settings = dict(
disable_heartbeat=False)
runner = CloudTaskRunner(task=Task())
runner.task_run_id = "foo"
res = runner._heartbeat()
assert res is True
assert runner.task_run_id == "foo"
assert runner.heartbeat_cmd == [
"prefect", "heartbeat", "task-run", "-i", "foo"
]
def test_task_runner_does_not_have_heartbeat_if_disabled(self, client):
# mock the returned value from the flow
mock = MagicMock(
flow_run=MagicMock(
flow=MagicMock(settings=MagicMock(disable_heartbeat=True))
)
)
client.graphql = MagicMock(return_value=mock)
# set up the CloudTaskRunner and its task
@prefect.task
def sleeper():
time.sleep(2)
runner = CloudTaskRunner(task=sleeper)
assert runner._heartbeat() is False
