diff --git a/CHANGELOG.md b/CHANGELOG.md index 9bd8829ff8e0..8a9a331827f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/ ### Enhancements +- Update Cloud config name for heartbeat settings - [#2081](https://github.com/PrefectHQ/prefect/pull/2081) - Add examples to Interactive API Docs - [#2122](https://github.com/PrefectHQ/prefect/pull/2122) - Use a new boto3 session per thread when using S3ResultHandlers- [#2108](https://github.com/PrefectHQ/prefect/issues/2108) - Allow users to skip Docker healthchecks - [#2150](https://github.com/PrefectHQ/prefect/pull/2150) diff --git a/src/prefect/engine/cloud/flow_runner.py b/src/prefect/engine/cloud/flow_runner.py index 9247c6cc65e7..8d5390dabee6 100644 --- a/src/prefect/engine/cloud/flow_runner.py +++ b/src/prefect/engine/cloud/flow_runner.py @@ -75,7 +75,7 @@ def _heartbeat(self) -> bool: } } flow_run = self.client.graphql(query).data.flow_run_by_pk - if flow_run.flow.settings.get("disable_heartbeat"): + if not flow_run.flow.settings.get("heartbeat_enabled", True): return False return True except Exception as exc: diff --git a/src/prefect/engine/cloud/task_runner.py b/src/prefect/engine/cloud/task_runner.py index 2f728ae2ae36..ada5820d034e 100644 --- a/src/prefect/engine/cloud/task_runner.py +++ b/src/prefect/engine/cloud/task_runner.py @@ -72,7 +72,7 @@ def _heartbeat(self) -> bool: } } flow_run = self.client.graphql(query).data.flow_run_by_pk - if flow_run.flow.settings.get("disable_heartbeat"): + if not flow_run.flow.settings.get("heartbeat_enabled", True): return False return True except Exception as exc: diff --git a/tests/engine/cloud/test_cloud_flow_runner.py b/tests/engine/cloud/test_cloud_flow_runner.py index 305bffb73476..ee48f7182a87 100644 --- a/tests/engine/cloud/test_cloud_flow_runner.py +++ b/tests/engine/cloud/test_cloud_flow_runner.py @@ -387,14 +387,17 @@ def test_heartbeat_traps_errors_caused_by_client(caplog, monkeypatch): assert "Heartbeat failed for Flow 'bad'" in log.message -def test_flow_runner_heartbeat_sets_command(monkeypatch): +@pytest.mark.parametrize("setting_available", [True, False]) +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( - disable_heartbeat=False + + 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() @@ -409,7 +412,7 @@ def test_flow_runner_does_not_have_heartbeat_if_disabled(monkeypatch): "prefect.engine.cloud.flow_runner.Client", MagicMock(return_value=client) ) client.graphql.return_value.data.flow_run_by_pk.flow.settings = dict( - disable_heartbeat=True + heartbeat_enabled=False ) # set up the CloudFlowRunner diff --git a/tests/engine/cloud/test_cloud_task_runner.py b/tests/engine/cloud/test_cloud_task_runner.py index 345dc3a5da59..0e4db8893b8b 100644 --- a/tests/engine/cloud/test_cloud_task_runner.py +++ b/tests/engine/cloud/test_cloud_task_runner.py @@ -483,14 +483,16 @@ def test_heartbeat_traps_errors_caused_by_bad_attributes(self, caplog, monkeypat assert log.levelname == "ERROR" assert "Heartbeat failed for Task 'Task'" in log.message - def test_task_runner_heartbeat_sets_command(self, monkeypatch): + @pytest.mark.parametrize("setting_available", [True, False]) + def test_task_runner_heartbeat_sets_command(self, monkeypatch, setting_available): 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 + client.graphql.return_value.data.flow_run_by_pk.flow.settings = ( + dict(heartbeat_enabled=True) if setting_available else {} ) + runner = CloudTaskRunner(task=Task()) runner.task_run_id = "foo" res = runner._heartbeat() @@ -504,7 +506,7 @@ def test_task_runner_does_not_have_heartbeat_if_disabled(self, monkeypatch): "prefect.engine.cloud.task_runner.Client", MagicMock(return_value=client) ) client.graphql.return_value.data.flow_run_by_pk.flow.settings = dict( - disable_heartbeat=True + heartbeat_enabled=False ) runner = CloudTaskRunner(task=Task()) runner.task_run_id = "foo"