feat: Add flagsmith healthcheck command

Author: khvn26

src/common/core/cli/__init__.py

@@ -0,0 +1,39 @@
+import argparse
+import socket
+
+
+def get_args(argv: list[str]) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Check if the API is able to accept local TCP connections.",
+    )
+    parser.add_argument(
+        "--port",
+        "-p",
+        type=int,
+        default=8000,
+        help="Port to check the API on (default: 8000)",
+    )
+    parser.add_argument(
+        "--timeout",
+        "-t",
+        type=int,
+        default=1,
+        help="Socket timeout for the connection attempt in seconds (default: 1)",
+    )
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str]) -> None:
+    args = get_args(argv)
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.settimeout(args.timeout)
+    try:
+        sock.connect(("localhost", args.port))
+    except socket.error as e:
+        sock.close()
+        print(f"Failed: {e} {args.port=}")
+        exit(1)
+    else:
+        sock.close()
+        print("OK")
+        exit(0)

src/common/core/cli/healthcheck.py

@@ -5,7 +5,11 @@
 import tempfile
 import typing
 
-from django.core.management import execute_from_command_line
+from django.core.management import (
+    execute_from_command_line as django_execute_from_command_line,
+)
+
+from common.core.cli import healthcheck
 
 logger = logging.getLogger(__name__)
 
@@ -56,7 +60,19 @@ def ensure_cli_env() -> typing.Generator[None, None, None]:
         yield
 
 
-def main() -> None:
+def execute_from_command_line(argv: list[str]) -> None:
+    try:
+        subcommand = argv[1]
+        subcommand_main = {
+            "healthcheck": healthcheck.main,
+        }[subcommand]
+    except (IndexError, KeyError):
+        django_execute_from_command_line(argv)
+    else:
+        subcommand_main(argv[2:])
+
+
+def main(argv: list[str] = sys.argv) -> None:
     """
     The main entry point to the Flagsmith application.
 
@@ -72,5 +88,5 @@ def main() -> None:
     `flagsmith <command> [options]`
     """
     with ensure_cli_env():
-        # Run Django
-        execute_from_command_line(sys.argv)
+        # Run own commands and Django
+        execute_from_command_line(argv)

src/common/core/main.py

@@ -0,0 +1,70 @@
+import io
+import subprocess
+
+import django
+import pytest
+from django.core.management import ManagementUtility
+
+from common.core.main import main
+
+
+@pytest.mark.parametrize(
+    "argv, expected_out",
+    (
+        [["flagsmith"], ManagementUtility().main_help_text()],
+        [["flagsmith", "version"], django.get_version()],
+    ),
+)
+def test_main__non_overridden_args__defaults_to_django(
+    argv: list[str],
+    expected_out: str,
+    capsys: pytest.CaptureFixture[str],
+) -> None:
+    # Given & When
+    main(argv)
+
+    # Then
+    assert capsys.readouterr().out[:-1] == expected_out
+
+
+def test_main__healthcheck__no_server__runs_expected(
+    capsys: pytest.CaptureFixture[str],
+) -> None:
+    # Given
+    argv = ["flagsmith", "healthcheck"]
+
+    # When
+    with pytest.raises(SystemExit) as exc_info:
+        main(argv)
+
+    # Then
+    assert exc_info.value.code == 1
+    assert (
+        capsys.readouterr().out[:-1]
+        == "Failed: [Errno 61] Connection refused args.port=8000"
+    )
+
+
+def test_main__healtcheck__server__runs_expected(
+    unused_tcp_port: int,
+    capsys: pytest.CaptureFixture[str],
+) -> None:
+    # Given
+    process = subprocess.Popen(
+        ["flagsmith", "start", "-b", f"0.0.0.0:{unused_tcp_port}", "api"],
+        stderr=subprocess.PIPE,
+    )
+    assert process.stderr
+    for line in io.TextIOWrapper(process.stderr, "utf-8"):
+        if "Listening" in line:
+            break
+
+    argv = ["flagsmith", "healthcheck", "--port", str(unused_tcp_port)]
+
+    # When
+    with pytest.raises(SystemExit) as exc_info:
+        main(argv)
+
+    # Then
+    assert exc_info.value.code == 0
+    assert capsys.readouterr().out[:-1] == "OK"