Merge pull request #151 from Flagsmith/feat/liveness-readiness

feat: Add liveness and readiness check endpoints. Return 503 instead of 500 when checks fail

Author: rolodato

src/edge_proxy/server.py

@@ -5,7 +5,7 @@
 from fastapi import FastAPI, Header
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
-from fastapi.responses import ORJSONResponse
+from fastapi.responses import ORJSONResponse, Response
 
 from edge_proxy.health_check.responses import HealthCheckResponse
 from fastapi_utils.tasks import repeat_every
@@ -40,11 +40,12 @@ async def unknown_key_error(request, exc):
 
 @app.get("/health", response_class=ORJSONResponse, deprecated=True)
 @app.get("/proxy/health", response_class=ORJSONResponse)
+@app.get("/proxy/health/readiness", response_class=ORJSONResponse)
 async def health_check():
     last_updated_at = environment_service.last_updated_at
     if not last_updated_at:
         return HealthCheckResponse(
-            status_code=500,
+            status_code=503,
             status="error",
             reason="environment document(s) not updated.",
             last_successful_update=None,
@@ -58,7 +59,7 @@ async def health_check():
         )
         if last_updated_at < threshold:
             return HealthCheckResponse(
-                status_code=500,
+                status_code=503,
                 status="error",
                 reason="environment document(s) stale.",
                 last_successful_update=last_updated_at,
@@ -67,6 +68,11 @@ async def health_check():
     return HealthCheckResponse(last_successful_update=last_updated_at)
 
 
+@app.get("/proxy/health/liveness")
+async def liveness_check():
+    return Response(status_code=200)
+
+
 @app.get("/api/v1/flags/", response_class=ORJSONResponse)
 async def flags(feature: str = None, x_environment_key: str = Header(None)):
     try:

tests/test_health.py

@@ -0,0 +1,91 @@
+from datetime import datetime, timedelta
+
+import pytest
+from fastapi.testclient import TestClient
+from pytest_mock import MockerFixture
+
+from edge_proxy.settings import HealthCheckSettings
+
+READINESS_ENDPOINTS = [
+    "/proxy/health/readiness",
+    "/proxy/health",
+    "/health",
+]
+
+
+def test_liveness_check(client: TestClient) -> None:
+    response = client.get("/proxy/health/liveness")
+    assert response.status_code == 200
+
+
+@pytest.mark.parametrize("endpoint", READINESS_ENDPOINTS)
+def test_health_check_returns_200_if_cache_was_updated_recently(
+    mocker: MockerFixture,
+    client: TestClient,
+    endpoint: str,
+) -> None:
+    mocked_environment_service = mocker.patch("edge_proxy.server.environment_service")
+    mocked_environment_service.last_updated_at = datetime.now()
+
+    response = client.get(endpoint)
+    assert response.status_code == 200
+
+
+@pytest.mark.parametrize("endpoint", READINESS_ENDPOINTS)
+def test_health_check_returns_503_if_cache_was_not_updated(
+    client: TestClient,
+    endpoint: str,
+) -> None:
+    response = client.get(endpoint)
+    assert response.status_code == 503
+    assert response.json() == {
+        "status": "error",
+        "reason": "environment document(s) not updated.",
+        "last_successful_update": None,
+    }
+
+
+@pytest.mark.parametrize("endpoint", READINESS_ENDPOINTS)
+def test_health_check_returns_503_if_cache_is_stale(
+    mocker: MockerFixture,
+    client: TestClient,
+    endpoint: str,
+) -> None:
+    last_updated_at = datetime.now() - timedelta(days=10)
+    mocked_environment_service = mocker.patch("edge_proxy.server.environment_service")
+    mocked_environment_service.last_updated_at = last_updated_at
+
+    response = client.get(endpoint)
+
+    assert response.status_code == 503
+    assert response.json() == {
+        "status": "error",
+        "reason": "environment document(s) stale.",
+        "last_successful_update": last_updated_at.isoformat(),
+    }
+
+
+@pytest.mark.parametrize("endpoint", READINESS_ENDPOINTS)
+def test_health_check_returns_200_if_cache_is_never_stale(
+    mocker: MockerFixture,
+    client: TestClient,
+    endpoint: str,
+) -> None:
+    # Given
+    health_check = HealthCheckSettings(environment_update_grace_period_seconds=None)
+    mocker.patch("edge_proxy.server.settings.health_check", health_check)
+
+    last_updated_at = datetime.now() - timedelta(days=10)
+    mocked_environment_service = mocker.patch("edge_proxy.server.environment_service")
+    mocked_environment_service.last_updated_at = last_updated_at
+
+    # When
+    response = client.get(endpoint)
+
+    # Then
+    assert response.status_code == 200
+    assert response.json() == {
+        "status": "ok",
+        "reason": None,
+        "last_successful_update": last_updated_at.isoformat(),
+    }

tests/test_server.py

@@ -1,83 +1,15 @@
-from datetime import datetime, timedelta
 import typing
 
 import orjson
-import pytest
 from fastapi.testclient import TestClient
 from pytest_mock import MockerFixture
 
-from edge_proxy.settings import HealthCheckSettings
 from tests.fixtures.response_data import environment_1
 
 if typing.TYPE_CHECKING:
     from edge_proxy.environments import EnvironmentService
 
 
-@pytest.mark.parametrize("endpoint", ["/proxy/health", "/health"])
-def test_health_check_returns_200_if_cache_was_updated_recently(
-    mocker: MockerFixture,
-    endpoint: str,
-    client: TestClient,
-) -> None:
-    mocked_environment_service = mocker.patch("edge_proxy.server.environment_service")
-    mocked_environment_service.last_updated_at = datetime.now()
-
-    response = client.get(endpoint)
-    assert response.status_code == 200
-
-
-def test_health_check_returns_500_if_cache_was_not_updated(
-    client: TestClient,
-) -> None:
-    response = client.get("/proxy/health")
-    assert response.status_code == 500
-    assert response.json() == {
-        "status": "error",
-        "reason": "environment document(s) not updated.",
-        "last_successful_update": None,
-    }
-
-
-def test_health_check_returns_500_if_cache_is_stale(
-    mocker: MockerFixture,
-    client: TestClient,
-) -> None:
-    last_updated_at = datetime.now() - timedelta(days=10)
-    mocked_environment_service = mocker.patch("edge_proxy.server.environment_service")
-    mocked_environment_service.last_updated_at = last_updated_at
-    response = client.get("/proxy/health")
-    assert response.status_code == 500
-    assert response.json() == {
-        "status": "error",
-        "reason": "environment document(s) stale.",
-        "last_successful_update": last_updated_at.isoformat(),
-    }
-
-
-def test_health_check_returns_200_if_cache_is_stale_and_health_check_configured_correctly(
-    mocker: MockerFixture,
-    client: TestClient,
-) -> None:
-    # Given
-    health_check = HealthCheckSettings(environment_update_grace_period_seconds=None)
-    mocker.patch("edge_proxy.server.settings.health_check", health_check)
-
-    last_updated_at = datetime.now() - timedelta(days=10)
-    mocked_environment_service = mocker.patch("edge_proxy.server.environment_service")
-    mocked_environment_service.last_updated_at = last_updated_at
-
-    # When
-    response = client.get("/proxy/health")
-
-    # Then
-    assert response.status_code == 200
-    assert response.json() == {
-        "status": "ok",
-        "reason": None,
-        "last_successful_update": last_updated_at.isoformat(),
-    }
-
-
 def test_get_flags(
     mocker: MockerFixture,
     environment_1_feature_states_response_list: list[dict],