From 1b948f4d5a5b13b81235b653bb16533e3d252290 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C4=93teris=20Caune?= <cuu508@gmail.com>
Date: Wed, 15 May 2019 13:42:38 +0300
Subject: [PATCH] Show check counts in JSON "badges". Fixes #251

---
 CHANGELOG.md               |  1 +
 hc/api/tests/test_badge.py | 26 +++++++++++++++++++++++---
 hc/api/views.py            | 22 +++++++++++++++++-----
 3 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 139451bc..493f50c9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@ All notable changes to this project will be documented in this file.
 
 ### Improvements
 - Add the `prunetokenbucket` management command
+- Show check counts in JSON "badges" (#251)
 
 ### Bug Fixes
 - Fix badges for tags containing special characters (#240, #237)
diff --git a/hc/api/tests/test_badge.py b/hc/api/tests/test_badge.py
index 038483b0..9b1dbeae 100644
--- a/hc/api/tests/test_badge.py
+++ b/hc/api/tests/test_badge.py
@@ -33,6 +33,14 @@ class BadgeTestCase(BaseTestCase):
         self.assertEqual(r.status_code, 204)
         self.assertEqual(r["Access-Control-Allow-Origin"], "*")
 
+    def test_it_handles_new(self):
+        r = self.client.get(self.json_url)
+        doc = r.json()
+        self.assertEqual(doc["status"], "up")
+        self.assertEqual(doc["total"], 1)
+        self.assertEqual(doc["grace"], 0)
+        self.assertEqual(doc["down"], 0)
+
     def test_it_handles_started_but_down(self):
         self.check.last_start = now()
         self.check.tags = "foo"
@@ -40,7 +48,11 @@ class BadgeTestCase(BaseTestCase):
         self.check.save()
 
         r = self.client.get(self.json_url)
-        self.assertContains(r, "down")
+        doc = r.json()
+        self.assertEqual(doc["status"], "down")
+        self.assertEqual(doc["total"], 1)
+        self.assertEqual(doc["grace"], 0)
+        self.assertEqual(doc["down"], 1)
 
     def test_it_shows_grace_badge(self):
         self.check.last_ping = now() - td(days=1, minutes=10)
@@ -49,7 +61,11 @@ class BadgeTestCase(BaseTestCase):
         self.check.save()
 
         r = self.client.get(self.json_url)
-        self.assertContains(r, "late")
+        doc = r.json()
+        self.assertEqual(doc["status"], "late")
+        self.assertEqual(doc["total"], 1)
+        self.assertEqual(doc["grace"], 1)
+        self.assertEqual(doc["down"], 0)
 
     def test_it_shows_started_but_grace_badge(self):
         self.check.last_start = now()
@@ -59,7 +75,11 @@ class BadgeTestCase(BaseTestCase):
         self.check.save()
 
         r = self.client.get(self.json_url)
-        self.assertContains(r, "late")
+        doc = r.json()
+        self.assertEqual(doc["status"], "late")
+        self.assertEqual(doc["total"], 1)
+        self.assertEqual(doc["grace"], 1)
+        self.assertEqual(doc["down"], 0)
 
     def test_it_handles_special_characters(self):
         self.check.tags = "db@dc1"
diff --git a/hc/api/views.py b/hc/api/views.py
index 526e588e..5f094483 100644
--- a/hc/api/views.py
+++ b/hc/api/views.py
@@ -200,7 +200,6 @@ def badge(request, badge_key, signature, tag, format="svg"):
     if not check_signature(badge_key, tag, signature):
         return HttpResponseNotFound()
 
-    status = "up"
     q = Check.objects.filter(project__badge_key=badge_key)
     if tag != "*":
         q = q.filter(tags__contains=tag)
@@ -208,20 +207,33 @@ def badge(request, badge_key, signature, tag, format="svg"):
     else:
         label = settings.MASTER_BADGE_LABEL
 
+    status, total, grace, down = "up", 0, 0, 0
     for check in q:
         if tag != "*" and tag not in check.tags_list():
             continue
 
+        total += 1
         check_status = check.get_status(with_started=False)
-        if status == "up" and check_status == "grace":
-            status = "late"
 
         if check_status == "down":
+            down += 1
             status = "down"
-            break
+            if format == "svg":
+                # For SVG badges, we can leave the loop as soon as we
+                # find the first "down"
+                break
+        elif check_status == "grace":
+            grace += 1
+            if status == "up":
+                status = "late"
 
     if format == "json":
-        return JsonResponse({"status": status})
+        return JsonResponse({
+            "status": status,
+            "total": total,
+            "grace": grace,
+            "down": down
+        })
 
     svg = get_badge_svg(label, status)
     return HttpResponse(svg, content_type="image/svg+xml")