From 52c1de370e5bade2ec7704074f26b042f2d9d9de Mon Sep 17 00:00:00 2001
From: Mathieu Leplatre <mathieu@mozilla.com>
Date: Tue, 3 Dec 2024 12:26:19 +0100
Subject: [PATCH] Fix #1422: add BigQuery check to heartbeat

---
 telescope/app.py              | 10 ++++++++++
 tests/conftest.py             |  7 +++++++
 tests/test_basic_endpoints.py |  4 +++-
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/telescope/app.py b/telescope/app.py
index 06161720..39812ab8 100644
--- a/telescope/app.py
+++ b/telescope/app.py
@@ -216,6 +216,7 @@ async def lbheartbeat(request):
 @routes.get("/__heartbeat__")
 async def heartbeat(request):
     checks = {}
+
     # Check that `curl` has HTTP2 and HTTP3 for `checks.core.http_versions`
     curl_cmd = subprocess.run(
         [config.CURL_BINARY_PATH, "--version"],
@@ -228,8 +229,17 @@ async def heartbeat(request):
         if not missing_features
         else f"missing features {', '.join(missing_features)}"
     )
+
+    # Bugzilla
     bz_ping = await request.app["telescope.tracker"].ping()
     checks["bugzilla"] = "ok" if bz_ping else "Bugzilla ping failed"
+
+    # Big Query
+    try:
+        checks["bigquery"] = (await utils.fetch_bigquery("SELECT 'ok';"))[0]
+    except Exception as exc:
+        checks["bigquery"] = str(exc)
+
     status = 200 if all(v == "ok" for v in checks.values()) else 503
     return web.json_response(checks, status=status)
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 0df1e74d..a530da57 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,5 +1,6 @@
 import os
 from typing import List, Union
+from unittest import mock
 
 import pytest
 import responses
@@ -58,6 +59,12 @@ def mock_aioresponses(cli):
         yield m
 
 
+@pytest.fixture
+def mock_bigquery_client():
+    with mock.patch("telescope.utils.bigquery.Client") as mocked:
+        yield mocked
+
+
 class ResponsesWrapper:
     """A tiny wrapper to mimic the aioresponses API."""
 
diff --git a/tests/test_basic_endpoints.py b/tests/test_basic_endpoints.py
index d39c7d32..8b9480cf 100644
--- a/tests/test_basic_endpoints.py
+++ b/tests/test_basic_endpoints.py
@@ -33,16 +33,18 @@ async def test_lbheartbeat(cli):
     assert response.status == 200
 
 
-async def test_heartbeat(cli, config, mock_aioresponses):
+async def test_heartbeat(cli, config, mock_aioresponses, mock_bigquery_client):
     config.BUGTRACKER_URL = "http://bugzilla.local"
     mock_aioresponses.get(
         config.BUGTRACKER_URL + "/rest/whoami", payload={"name": "foo"}
     )
+    mock_bigquery_client.return_value.query.side_effect = ValueError("bad credentials")
 
     response = await cli.get("/__heartbeat__")
     body = await response.json()
 
     assert body["bugzilla"] == "ok"
+    assert body["bigquery"] == "bad credentials"
     assert body["curl"] == "ok"
     assert response.status == 200