def test_flow_runner_does_not_have_heartbeat_if_disabled(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 CloudFlowRunner
runner = CloudFlowRunner(flow=prefect.Flow(name="test"))
# confirm the runner's heartbeat respects the heartbeat toggle
assert runner._heartbeat() is False
def test_flow_runner_does_not_have_heartbeat_if_disabled(monkeypatch):
client = MagicMock()
monkeypatch.setattr("prefect.engine.cloud.flow_runner.Client",
MagicMock(return_value=client))
client.graphql.return_value.data.flow_run_by_pk.flow.settings = dict(
heartbeat_enabled=False)
# set up the CloudFlowRunner
runner = CloudFlowRunner(flow=prefect.Flow(name="test"))
# confirm the runner's heartbeat respects the heartbeat toggle
assert runner._heartbeat() is False
def test_heartbeat_traps_errors_caused_by_client(caplog, monkeypatch):
client = MagicMock(graphql=MagicMock(side_effect=SyntaxError))
monkeypatch.setattr("prefect.engine.cloud.flow_runner.Client",
MagicMock(return_value=client))
runner = CloudFlowRunner(flow=prefect.Flow(name="bad"))
res = runner._heartbeat()
assert res is False
log = caplog.records[0]
assert log.levelname == "ERROR"
assert "Heartbeat failed for Flow 'bad'" in log.message
def test_heartbeat_traps_errors_caused_by_client(monkeypatch):
client = MagicMock(update_flow_run_heartbeat=MagicMock(side_effect=SyntaxError))
monkeypatch.setattr(
"prefect.engine.cloud.flow_runner.Client", MagicMock(return_value=client)
)
runner = CloudFlowRunner(flow=prefect.Flow(name="bad"))
with pytest.warns(UserWarning) as warning:
res = runner._heartbeat()
assert res is None
assert client.update_flow_run_heartbeat.called
w = warning.pop()
assert "Heartbeat failed for Flow 'bad'" in repr(w.message)
def test_flow_runner_heartbeat_sets_command(monkeypatch):
client = MagicMock()
monkeypatch.setattr("prefect.engine.cloud.flow_runner.Client",
MagicMock(return_value=client))
client.graphql.return_value.data.flow_run_by_pk.flow.settings = dict(
disable_heartbeat=False)
runner = CloudFlowRunner(flow=prefect.Flow(name="test"))
with prefect.context(flow_run_id="foo"):
res = runner._heartbeat()
assert res is True
assert runner.heartbeat_cmd == [
"prefect", "heartbeat", "flow-run", "-i", "foo"
]
def test_flow_runner_heartbeat_sets_command(monkeypatch, setting_available):
client = MagicMock()
monkeypatch.setattr("prefect.engine.cloud.flow_runner.Client",
MagicMock(return_value=client))
client.graphql.return_value.data.flow_run_by_pk.flow.settings = (dict(
heartbeat_enabled=True) if setting_available else {})
runner = CloudFlowRunner(flow=prefect.Flow(name="test"))
with prefect.context(flow_run_id="foo"):
res = runner._heartbeat()
assert res is True
assert runner.heartbeat_cmd == [
sys.executable,
"-m",
"prefect",
"heartbeat",
"flow-run",
"-i",
"foo",
]
