add http healthcheck, subcommands

Author: khvn26

src/common/core/cli/healthcheck.py

@@ -1,38 +1,99 @@
 import argparse
 import socket
+import urllib.parse
+
+import requests
 
 
 def get_args(argv: list[str]) -> argparse.Namespace:
     parser = argparse.ArgumentParser(
-        description="Check if the API is able to accept local TCP connections.",
+        description="Health checks",
+    )
+    subcommands = parser.add_subparsers(dest="subcommand", required=True)
+    tcp_parser = subcommands.add_parser(
+        "tcp", help="Check if the API is able to accept local TCP connections"
     )
-    parser.add_argument(
+    tcp_parser.add_argument(
         "--port",
         "-p",
         type=int,
         default=8000,
         help="Port to check the API on (default: 8000)",
     )
-    parser.add_argument(
+    tcp_parser.add_argument(
         "--timeout",
         "-t",
         type=int,
         default=1,
         help="Socket timeout for the connection attempt in seconds (default: 1)",
     )
+    http_parser = subcommands.add_parser(
+        "http", help="Check if the API is able to serve HTTP requests"
+    )
+    http_parser.add_argument(
+        "--port",
+        "-p",
+        type=int,
+        default=8000,
+        help="Port to check the API on (default: 8000)",
+    )
+    http_parser.add_argument(
+        "--timeout",
+        "-t",
+        type=int,
+        default=1,
+        help="Request timeout in seconds (default: 1)",
+    )
+    http_parser.add_argument(
+        "path",
+        nargs="?",
+        type=str,
+        default="/health/liveness",
+        help="Request path (default: /health/liveness)",
+    )
     return parser.parse_args(argv)
 
 
-def main(argv: list[str]) -> None:
-    args = get_args(argv)
+def check_tcp_connection(
+    port: int,
+    timeout_seconds: int,
+) -> None:
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    sock.settimeout(args.timeout)
+    sock.settimeout(timeout_seconds)
     try:
-        sock.connect(("127.0.0.1", args.port))
+        sock.connect(("127.0.0.1", port))
     except socket.error as e:
-        print(f"Failed: {e} {args.port=}")
+        print(f"Failed: {e} {port=}")
         exit(1)
     else:
         exit(0)
     finally:
         sock.close()
+
+
+def check_http_response(
+    port: int,
+    timeout_seconds: int,
+    path: str,
+) -> None:
+    url = urllib.parse.urljoin(f"http://127.0.0.1:{port}", path)
+    requests.get(
+        url,
+        timeout=timeout_seconds,
+    ).raise_for_status()
+
+
+def main(argv: list[str]) -> None:
+    args = get_args(argv)
+    match args.subcommand:
+        case "tcp":
+            check_tcp_connection(
+                port=args.port,
+                timeout_seconds=args.timeout,
+            )
+        case "http":
+            check_http_response(
+                port=args.port,
+                timeout_seconds=args.timeout,
+                path=args.path,
+            )

tests/integration/core/test_main.py

@@ -25,32 +25,64 @@ def test_main__non_overridden_args__defaults_to_django(
     assert capsys.readouterr().out[:-1] == expected_out
 
 
-def test_main__healthcheck__no_server__runs_expected(
+def test_main__healthcheck_tcp__no_server__runs_expected(
     capsys: pytest.CaptureFixture[str],
 ) -> None:
     # Given
-    argv = ["flagsmith", "healthcheck"]
+    argv = ["flagsmith", "healthcheck", "tcp"]
 
     # When
     with pytest.raises(SystemExit) as exc_info:
         main(argv)
 
     # Then
-    assert "Connection refused args.port=8000" in capsys.readouterr().out
+    assert "Connection refused port=8000" in capsys.readouterr().out
     assert exc_info.value.code == 1
 
 
-def test_main__healtcheck__server__runs_expected(
+def test_main__healtcheck_tcp__server__runs_expected(
     unused_tcp_port: int,
     http_server: HTTPServer,
-    capsys: pytest.CaptureFixture[str],
 ) -> None:
     # Given
-    argv = ["flagsmith", "healthcheck", "--port", str(unused_tcp_port)]
+    argv = ["flagsmith", "healthcheck", "tcp", "--port", str(unused_tcp_port)]
 
     # When
     with pytest.raises(SystemExit) as exc_info:
         main(argv)
 
     # Then
     assert exc_info.value.code == 0
+
+
+def test_main__healtcheck_http__no_server__runs_expected(
+    capsys: pytest.CaptureFixture[str],
+) -> None:
+    # Given
+    argv = ["flagsmith", "healthcheck", "http"]
+
+    # When & Then
+    with pytest.raises(Exception):
+        main(argv)
+
+
+@pytest.mark.parametrize(
+    "argv, expected_path",
+    (
+        (["flagsmith", "healthcheck", "http"], "/health/liveness"),
+        (["flagsmith", "healthcheck", "http", "health/readiness"], "/health/readiness"),
+    ),
+)
+def test_main__healthcheck_http__server__runs_expected(
+    unused_tcp_port: int,
+    http_server: HTTPServer,
+    argv: list[str],
+    expected_path: str,
+) -> None:
+    # Given
+    argv.extend(["--port", str(unused_tcp_port)])
+
+    http_server.expect_request(expected_path).respond_with_data(status=200)
+
+    # When & Then
+    main(argv)