From 34786181939126d7528c768aafbebc838c90c2b1 Mon Sep 17 00:00:00 2001
From: masayuki038 <masayuki038@gmail.com>
Date: Sat, 9 Dec 2023 15:21:58 +0900
Subject: [PATCH] Show an error message on "Test Connection" failure for Google
 Spreadsheet Query Runner

---
 redash/query_runner/google_spreadsheets.py    |  7 +++-
 .../query_runner/test_google_spreadsheets.py  | 35 ++++++++++++++++++-
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/redash/query_runner/google_spreadsheets.py b/redash/query_runner/google_spreadsheets.py
index 279a329dc1..17df3593f8 100644
--- a/redash/query_runner/google_spreadsheets.py
+++ b/redash/query_runner/google_spreadsheets.py
@@ -23,6 +23,7 @@
 try:
     import google.auth
     import gspread
+    from google.auth.exceptions import GoogleAuthError
     from google.oauth2.service_account import Credentials
     from gspread.exceptions import APIError
     from gspread.exceptions import WorksheetNotFound as GSWorksheetNotFound
@@ -230,13 +231,17 @@ def _get_spreadsheet_service(self):
         return spreadsheetservice
 
     def test_connection(self):
-        service = self._get_spreadsheet_service()
         test_spreadsheet_key = "1S0mld7LMbUad8LYlo13Os9f7eNjw57MqVC0YiCd1Jis"
         try:
+            service = self._get_spreadsheet_service()
             service.open_by_key(test_spreadsheet_key).worksheets()
         except APIError as e:
+            logger.exception(e)
             message = parse_api_error(e)
             raise Exception(message)
+        except GoogleAuthError as e:
+            logger.exception(e)
+            raise Exception(str(e))
 
     def run_query(self, query, user):
         logger.debug("Spreadsheet is about to execute query: %s", query)
diff --git a/tests/query_runner/test_google_spreadsheets.py b/tests/query_runner/test_google_spreadsheets.py
index 6f02e10e1a..db16d8fc3d 100644
--- a/tests/query_runner/test_google_spreadsheets.py
+++ b/tests/query_runner/test_google_spreadsheets.py
@@ -1,12 +1,16 @@
 import datetime
 from unittest import TestCase
 
-from mock import MagicMock
+import pytest
+from google.auth.exceptions import TransportError
+from gspread.exceptions import APIError
+from mock import MagicMock, patch
 
 from redash.query_runner import TYPE_DATETIME, TYPE_FLOAT
 from redash.query_runner.google_spreadsheets import (
     TYPE_BOOLEAN,
     TYPE_STRING,
+    GoogleSpreadsheet,
     WorksheetNotFoundByTitleError,
     WorksheetNotFoundError,
     _get_columns_and_column_names,
@@ -171,3 +175,32 @@ def test_is_url_key(self):
 
         _key = "key|0"
         self.assertFalse(is_url_key(_key))
+
+
+class TestConnection(TestCase):
+    @patch("redash.query_runner.google_spreadsheets.google.auth.default")
+    @patch("redash.query_runner.google_spreadsheets.gspread.Client")
+    def test_connect_succuess(self, mock_client, _mock_auth_default):
+        try:
+            qr_gspread = GoogleSpreadsheet({})
+            qr_gspread.test_connection()
+            mock_client().login.assert_called_once_with()
+            mock_client().open_by_key.assert_called_once()
+        except Exception:
+            self.fail("test_connection failed")
+
+    @patch("redash.query_runner.google_spreadsheets.google.auth.default")
+    def test_connect_fail_with_transport_error(self, mock_auth_default):
+        mock_auth_default.side_effect = TransportError("Connection Refused")
+        qr_gspread = GoogleSpreadsheet({})
+        with pytest.raises(Exception):
+            qr_gspread.test_connection()
+
+    @patch("redash.query_runner.google_spreadsheets.google.auth.default")
+    def test_connect_fail_with_api_error(self, mock_auth_default):
+        mock_response = MagicMock()
+        mock_response.json.return_value = {"error": {"message": "Sheet API is disabled"}}
+        mock_auth_default.side_effect = APIError(mock_response)
+        qr_gspread = GoogleSpreadsheet({})
+        with pytest.raises(Exception):
+            qr_gspread.test_connection()