Show an error message on "Test Connection" failure for Google Spreadsheet Query Runner (#6652)

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 @@ class TestIsUrlKey(TestCase):
 
         _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()