feat(client): Add drain command for the worker preStop hook

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()

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)