diff --git a/clients/python/src/taskbroker_client/drain.py b/clients/python/src/taskbroker_client/drain.py new file mode 100644 index 00000000..4773c740 --- /dev/null +++ b/clients/python/src/taskbroker_client/drain.py @@ -0,0 +1,28 @@ +import argparse +import time +from pathlib import Path + +from taskbroker_client.constants import DEFAULT_WORKER_DRAIN_FILE_PATH + + +def drain(seconds: float, path: str = DEFAULT_WORKER_DRAIN_FILE_PATH) -> None: + """ + Create the drain file, then sleep. Runs as the worker pod's preStop hook, + which must finish before Kubernetes sends SIGTERM. + """ + Path(path).touch() + time.sleep(seconds) + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser( + prog="python3 -m taskbroker_client.drain", + description="Create the worker drain file, then sleep.", + ) + parser.add_argument("seconds", type=float, help="how long to sleep after creating the file") + args = parser.parse_args(argv) + drain(args.seconds) + + +if __name__ == "__main__": + main() diff --git a/clients/python/tests/test_drain.py b/clients/python/tests/test_drain.py new file mode 100644 index 00000000..3a92a6a9 --- /dev/null +++ b/clients/python/tests/test_drain.py @@ -0,0 +1,21 @@ +from pathlib import Path +from unittest import mock + +from taskbroker_client.drain import drain, main + + +def test_drain_creates_file_before_sleeping(tmp_path: Path) -> None: + path = tmp_path / "draining" + + def sleep(seconds: float) -> None: + assert path.exists() + + with mock.patch("taskbroker_client.drain.time.sleep", side_effect=sleep) as sleep_mock: + drain(15, str(path)) + sleep_mock.assert_called_once_with(15) + + +def test_main_parses_seconds() -> None: + with mock.patch("taskbroker_client.drain.drain") as drain_mock: + main(["15"]) + drain_mock.assert_called_once_with(15.0)