Skip to content
This repository has been archived by the owner on Feb 22, 2023. It is now read-only.

Commit

Permalink
Add database connectivity to healthcheck endpoint (#1083)
Browse files Browse the repository at this point in the history
  • Loading branch information
krysal authored Feb 1, 2023
1 parent 624df07 commit f0e4395
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 4 deletions.
23 changes: 19 additions & 4 deletions api/catalog/api/views/health_views.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from django.conf import settings
from django.db import connection
from rest_framework import status
from rest_framework.exceptions import APIException
from rest_framework.request import Request
Expand All @@ -21,19 +22,33 @@ class HealthCheck(APIView):

swagger_schema = None

def _check_es(self) -> Response | None:
"""Check ES cluster health and raise an exception if ES is not healthy."""
@staticmethod
def _check_db() -> None:
"""
Check that the database is available.
Returns nothing if everything is OK, throws error otherwise.
"""
connection.ensure_connection()

@staticmethod
def _check_es() -> None:
"""
Check Elasticsearch cluster health.
Raises an exception if ES is not healthy.
"""
es_health = settings.ES.cluster.health(timeout="5s")

if es_health["timed_out"]:
raise ElasticsearchHealthcheckException("es_timed_out")

if (status := es_health["status"]) != "green":
raise ElasticsearchHealthcheckException(f"es_status_{status}")
if (es_status := es_health["status"]) != "green":
raise ElasticsearchHealthcheckException(f"es_status_{es_status}")

def get(self, request: Request):
if "check_es" in request.query_params:
self._check_es()
self._check_db()

return Response({"status": "200 OK"}, status=200)
13 changes: 13 additions & 0 deletions api/test/unit/views/health_views_test.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from unittest import mock

import pook
import pytest

Expand All @@ -16,11 +18,21 @@ def mock_health_response(status="green", timed_out=False):
)


@pytest.mark.django_db
def test_health_check_plain(api_client):
res = api_client.get("/healthcheck/")
assert res.status_code == 200


def test_health_check_calls__check_db(api_client):
with mock.patch(
"catalog.api.views.health_views.HealthCheck._check_db"
) as mock_check_db:
res = api_client.get("/healthcheck/")
assert res.status_code == 200
mock_check_db.assert_called_once()


def test_health_check_es_timed_out(api_client):
mock_health_response(timed_out=True)
pook.on()
Expand All @@ -42,6 +54,7 @@ def test_health_check_es_status_bad(status, api_client):
assert res.json()["detail"] == f"es_status_{status}"


@pytest.mark.django_db
def test_health_check_es_all_good(api_client):
mock_health_response(status="green")
pook.on()
Expand Down

0 comments on commit f0e4395

Please sign in to comment.