diff --git a/clients/python/src/taskbroker_client/healthcheck.py b/clients/python/src/taskbroker_client/healthcheck.py new file mode 100644 index 00000000..5465b559 --- /dev/null +++ b/clients/python/src/taskbroker_client/healthcheck.py @@ -0,0 +1,28 @@ +import argparse +import os +import sys + + +def healthcheck(path: str) -> None: + """ + Remove the health file, so the next check fails unless the worker touches it + again. Runs as the worker pod's liveness probe. + """ + try: + os.remove(path) + except FileNotFoundError: + sys.exit(f"{path} is missing: the worker has not touched it since the last probe") + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser( + prog="python3 -m taskbroker_client.healthcheck", + description="Remove the worker health file, or exit 1 if it is missing.", + ) + parser.add_argument("path", help="the health check file the worker touches") + args = parser.parse_args(argv) + healthcheck(args.path) + + +if __name__ == "__main__": + main() diff --git a/clients/python/tests/test_healthcheck.py b/clients/python/tests/test_healthcheck.py new file mode 100644 index 00000000..816b5518 --- /dev/null +++ b/clients/python/tests/test_healthcheck.py @@ -0,0 +1,19 @@ +from pathlib import Path + +import pytest + +from taskbroker_client.healthcheck import main + + +def test_fails_when_worker_has_not_touched_file_since_last_check(tmp_path: Path) -> None: + path = tmp_path / "health" + path.touch() + + main([str(path)]) + assert not path.exists() + + with pytest.raises(SystemExit) as exc: + main([str(path)]) + assert exc.value.code == ( + f"{path} is missing: the worker has not touched it since the last probe" + )