From fdc014891b9a00deccba76856c9dae2e06c2de07 Mon Sep 17 00:00:00 2001
From: Graham Beckley <gbeckley@mozilla.com>
Date: Thu, 10 Aug 2023 14:15:50 -0400
Subject: [PATCH] Use centralized check runner in Django

---
 src/dockerflow/django/views.py | 52 ++++++++++++----------------------
 tests/django/test_django.py    | 50 +++++++++++++++++++++++++++-----
 2 files changed, 61 insertions(+), 41 deletions(-)

diff --git a/src/dockerflow/django/views.py b/src/dockerflow/django/views.py
index 3cc16e3..741bed5 100644
--- a/src/dockerflow/django/views.py
+++ b/src/dockerflow/django/views.py
@@ -2,11 +2,12 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, you can obtain one at http://mozilla.org/MPL/2.0/.
 from django.conf import settings
-from django.core import checks
+from django.core.checks.registry import registry as django_check_registry
 from django.http import HttpResponse, HttpResponseNotFound, JsonResponse
 from django.utils.module_loading import import_string
 
-from .checks import level_to_text
+from dockerflow import checks
+
 from .signals import heartbeat_failed, heartbeat_passed
 
 version_callback = getattr(
@@ -42,42 +43,25 @@ def heartbeat(request):
     Any check that returns a warning or worse (error, critical) will
     return a 500 response.
     """
-    all_checks = checks.registry.registry.get_checks(
-        include_deployment_checks=not settings.DEBUG
+    checks_to_run = (
+        (check.__name__, lambda: check(app_configs=None))
+        for check in django_check_registry.get_checks(
+            include_deployment_checks=not settings.DEBUG
+        )
     )
-
-    details = {}
-    statuses = {}
-    level = 0
-
-    for check in all_checks:
-        detail = heartbeat_check_detail(check)
-        statuses[check.__name__] = detail["status"]
-        level = max(level, detail["level"])
-        if detail["level"] > 0:
-            details[check.__name__] = detail
-
-    if level < checks.messages.ERROR:
+    check_results = checks.run_checks(
+        checks_to_run,
+        silenced_check_ids=settings.SILENCED_SYSTEM_CHECKS,
+    )
+    if check_results.level < checks.ERROR:
         status_code = 200
-        heartbeat_passed.send(sender=heartbeat, level=level)
+        heartbeat_passed.send(sender=heartbeat, level=check_results.level)
     else:
         status_code = 500
-        heartbeat_failed.send(sender=heartbeat, level=level)
+        heartbeat_failed.send(sender=heartbeat, level=check_results.level)
 
-    payload = {"status": level_to_text(level)}
+    payload = {"status": checks.level_to_text(check_results.level)}
     if settings.DEBUG:
-        payload["checks"] = statuses
-        payload["details"] = details
+        payload["checks"] = check_results.statuses
+        payload["details"] = check_results.details
     return JsonResponse(payload, status=status_code)
-
-
-def heartbeat_check_detail(check):
-    errors = check(app_configs=None)
-    errors = list(filter(lambda e: e.id not in settings.SILENCED_SYSTEM_CHECKS, errors))
-    level = max([0] + [e.level for e in errors])
-
-    return {
-        "status": level_to_text(level),
-        "level": level,
-        "messages": {e.id: e.msg for e in errors},
-    }
diff --git a/tests/django/test_django.py b/tests/django/test_django.py
index 9452a85..28e63ed 100644
--- a/tests/django/test_django.py
+++ b/tests/django/test_django.py
@@ -50,9 +50,8 @@ def test_version_missing(dockerflow_middleware, mocker, rf):
 
 
 @pytest.mark.django_db
-def test_heartbeat(dockerflow_middleware, reset_checks, rf, settings):
-    request = rf.get("/__heartbeat__")
-    response = dockerflow_middleware.process_request(request)
+def test_heartbeat(client, settings):
+    response = client.get("/__heartbeat__")
     assert response.status_code == 200
 
     settings.DOCKERFLOW_CHECKS = [
@@ -60,8 +59,46 @@ def test_heartbeat(dockerflow_middleware, reset_checks, rf, settings):
         "tests.django.django_checks.error",
     ]
     checks.register()
-    response = dockerflow_middleware.process_request(request)
+    response = client.get("/__heartbeat__")
+    assert response.status_code == 500
+    content = response.json()
+    assert content["status"] == "error"
+    assert content.get("checks") is None
+    assert content.get("details") is None
+
+
+@pytest.mark.django_db
+def test_heartbeat_debug(client, settings):
+    settings.DOCKERFLOW_CHECKS = [
+        "tests.django.django_checks.warning",
+        "tests.django.django_checks.error",
+    ]
+    settings.DEBUG = True
+    checks.register()
+    response = client.get("/__heartbeat__")
     assert response.status_code == 500
+    content = response.json()
+    assert content["status"]
+    assert content["checks"]
+    assert content["details"]
+
+
+@pytest.mark.django_db
+def test_heartbeat_silenced(client, settings):
+    settings.DOCKERFLOW_CHECKS = [
+        "tests.django.django_checks.warning",
+        "tests.django.django_checks.error",
+    ]
+    settings.SILENCED_SYSTEM_CHECKS.append("tests.checks.E001")
+    settings.DEBUG = True
+    checks.register()
+
+    response = client.get("/__heartbeat__")
+    assert response.status_code == 200
+    content = response.json()
+    assert content["status"] == "warning"
+    assert "warning" in content["details"]
+    assert "error" not in content["details"]
 
 
 @pytest.mark.django_db
@@ -75,11 +112,10 @@ def test_lbheartbeat_makes_no_db_queries(dockerflow_middleware, rf):
 
 
 @pytest.mark.django_db
-def test_redis_check(dockerflow_middleware, reset_checks, rf, settings):
+def test_redis_check(client, settings):
     settings.DOCKERFLOW_CHECKS = ["dockerflow.django.checks.check_redis_connected"]
     checks.register()
-    request = rf.get("/__heartbeat__")
-    response = dockerflow_middleware.process_request(request)
+    response = client.get("/__heartbeat__")
     assert response.status_code == 200