Andrushika commented on code in PR #73806:
URL: https://github.com/apache/airflow/pull/73806#discussion_r4204468297
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -1305,6 +1306,204 @@ def test_overtime_handling(
mock_kill.assert_not_called()
mock_logger.warning.assert_not_called()
+ def test_server_terminated_task_gets_killed_task_cleanup_time(self,
mocker, monkeypatch):
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
42.0)
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.MIN_HEARTBEAT_INTERVAL",
0)
+ mock_kill =
mocker.patch("airflow.sdk.execution_time.supervisor.WatchedSubprocess.kill")
+ client = mocker.Mock()
+ client.task_instances.heartbeat.side_effect =
ServerResponseError.from_response(
+ httpx.Response(
+ 409,
+ request=httpx.Request("PUT", "http://server/heartbeat";),
+ json={"detail": {"reason": "not_running", "current_state":
"failed"}},
+ )
+ )
+ proc = ActivitySubprocess(
+ process_log=mocker.MagicMock(),
+ id=TI_ID,
+ pid=12345,
+ stdin=mocker.Mock(),
+ process=mocker.Mock(),
+ client=client,
+ )
+
+ proc._send_heartbeat_if_needed()
+
+ mock_kill.assert_called_once_with(signal.SIGTERM, force=True,
escalation_delay=42.0)
+
+ def test_set_execution_timeout_schedules_enforcement(self, mocker,
monkeypatch):
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
5.0)
+ mocker.patch("time.monotonic", return_value=100.0)
Review Comment:
Freezegun seems not to be in our dependencies -- so I tried `time_machine`.
But it doesn't mock `time.monotonic`.
So I went with an advancing `time.monotonic` mock like the overtime test
below does. (Please lmk if freezegun is still preferred)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
