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", ]