diff --git a/arroyo/healthcheck.py b/arroyo/healthcheck.py new file mode 100644 index 00000000..a2cb56e2 --- /dev/null +++ b/arroyo/healthcheck.py @@ -0,0 +1,31 @@ +import argparse +import os +import sys +from typing import Optional + + +def healthcheck(path: str) -> None: + """ + Remove the file the Healthcheck strategy touches, so the next check fails + unless the consumer touches it again. Runs as the consumer's liveness probe. + """ + try: + os.remove(path) + except FileNotFoundError: + sys.exit( + f"{path} is missing: the consumer has not touched it since the last probe" + ) + + +def main(argv: Optional[list[str]] = None) -> None: + parser = argparse.ArgumentParser( + prog="python3 -m arroyo.healthcheck", + description="Remove the consumer health file, or exit 1 if it is missing.", + ) + parser.add_argument("path", help="the file the Healthcheck strategy touches") + args = parser.parse_args(argv) + healthcheck(args.path) + + +if __name__ == "__main__": + main() diff --git a/tests/test_healthcheck.py b/tests/test_healthcheck.py new file mode 100644 index 00000000..402a0d4f --- /dev/null +++ b/tests/test_healthcheck.py @@ -0,0 +1,23 @@ +from pathlib import Path +from unittest import mock + +import pytest + +from arroyo.healthcheck import main +from arroyo.processing.strategies.healthcheck import Healthcheck + + +def test_fails_when_consumer_has_not_touched_file_since_last_check( + tmp_path: Path, +) -> None: + path = tmp_path / "health.txt" + Healthcheck(str(path), mock.Mock()).poll() + + 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 consumer has not touched it since the last probe" + )