From 4727cb7893df6350a853fdf5cde8542c9c18b51a Mon Sep 17 00:00:00 2001
From: Max Risuhin <risuhin.max@gmail.com>
Date: Sat, 28 Dec 2019 05:22:37 -0800
Subject: [PATCH] Improve Google Drive auth failure error reporting

---
 dvc/remote/gdrive.py             | 43 ++++++++++++++++++++++-----
 tests/unit/remote/test_gdrive.py | 51 +++++++++++++++++++++++++++-----
 2 files changed, 80 insertions(+), 14 deletions(-)

diff --git a/dvc/remote/gdrive.py b/dvc/remote/gdrive.py
index 75251c9d2b..38fc819c83 100644
--- a/dvc/remote/gdrive.py
+++ b/dvc/remote/gdrive.py
@@ -21,8 +21,18 @@
 
 
 class GDriveRetriableError(DvcException):
-    def __init__(self, msg):
-        super(GDriveRetriableError, self).__init__(msg)
+    def __init__(self, msg, cause=None):
+        super(GDriveRetriableError, self).__init__(msg, cause=cause)
+
+
+class GDriveAccessTokenRefreshError(DvcException):
+    def __init__(self, msg, cause=None):
+        super(GDriveAccessTokenRefreshError, self).__init__(msg, cause=cause)
+
+
+class GDriveMissedCredentialKeyError(DvcException):
+    def __init__(self, msg, cause=None):
+        super(GDriveMissedCredentialKeyError, self).__init__(msg, cause=cause)
 
 
 @decorator
@@ -38,7 +48,7 @@ def _wrap_pydrive_retriable(call):
             "HttpError {}".format(code) in str(exception)
             for code in retry_codes
         ):
-            raise GDriveRetriableError(msg="Google API request failed")
+            raise GDriveRetriableError("Google API request failed")
         raise
     return result
 
@@ -78,7 +88,7 @@ def init_drive(self):
                 "https://man.dvc.org/remote/add."
             )
         self.gdrive_user_credentials_path = (
-            tmp_fname(".dvc/tmp/")
+            tmp_fname(os.path.join(self.repo.tmp_dir, ""))
             if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA)
             else self.config.get(
                 Config.SECTION_GDRIVE_USER_CREDENTIALS_FILE,
@@ -161,6 +171,8 @@ def cached_ids(self):
     @property
     @wrap_with(threading.RLock())
     def drive(self):
+        from pydrive.auth import RefreshError
+
         if not hasattr(self, "_gdrive"):
             from pydrive.auth import GoogleAuth
             from pydrive.drive import GoogleDrive
@@ -195,10 +207,27 @@ def drive(self):
 
             # Pass non existent settings path to force DEFAULT_SETTINGS loading
             gauth = GoogleAuth(settings_file="")
-            gauth.CommandLineAuth()
 
-            if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA):
-                os.remove(self.gdrive_user_credentials_path)
+            try:
+                gauth.CommandLineAuth()
+            except RefreshError as e:
+                raise GDriveAccessTokenRefreshError(
+                    "Google Drive's access token refreshment is failed", e
+                )
+            except KeyError as e:
+                raise GDriveMissedCredentialKeyError(
+                    "Google Drive's user credentials file '{}' "
+                    "misses value for key '{}'".format(
+                        self.gdrive_user_credentials_path, str(e)
+                    ),
+                    e,
+                )
+            # Handle pydrive.auth.AuthenticationError and others auth failures
+            except Exception as e:
+                raise DvcException("Google Drive authentication failed", e)
+            finally:
+                if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA):
+                    os.remove(self.gdrive_user_credentials_path)
 
             self._gdrive = GoogleDrive(gauth)
 
diff --git a/tests/unit/remote/test_gdrive.py b/tests/unit/remote/test_gdrive.py
index a587002781..a886fb6dd2 100644
--- a/tests/unit/remote/test_gdrive.py
+++ b/tests/unit/remote/test_gdrive.py
@@ -1,9 +1,46 @@
-import mock
-from dvc.remote.gdrive import RemoteGDrive
+from unittest import TestCase
 
+import pytest
+import os
 
-@mock.patch("dvc.remote.gdrive.RemoteGDrive.init_drive")
-def test_init(repo):
-    url = "gdrive://root/data"
-    gdrive = RemoteGDrive(repo, {"url": url})
-    assert str(gdrive.path_info) == url
+from dvc.remote.gdrive import (
+    RemoteGDrive,
+    GDriveAccessTokenRefreshError,
+    GDriveMissedCredentialKeyError,
+)
+
+USER_CREDS_TOKEN_REFRESH_ERROR = '{"access_token": "", "client_id": "", "client_secret": "", "refresh_token": "", "token_expiry": "", "token_uri": "https://oauth2.googleapis.com/token", "user_agent": null, "revoke_uri": "https://oauth2.googleapis.com/revoke", "id_token": null, "id_token_jwt": null, "token_response": {"access_token": "", "expires_in": 3600, "scope": "https://www.googleapis.com/auth/drive.appdata https://www.googleapis.com/auth/drive", "token_type": "Bearer"}, "scopes": ["https://www.googleapis.com/auth/drive", "https://www.googleapis.com/auth/drive.appdata"], "token_info_uri": "https://oauth2.googleapis.com/tokeninfo", "invalid": true, "_class": "OAuth2Credentials", "_module": "oauth2client.client"}'  # noqa: E501
+
+USER_CREDS_MISSED_KEY_ERROR = "{}"
+
+
+class Repo(object):
+    tmp_dir = ""
+
+
+class TestRemoteGDrive(TestCase):
+    CONFIG = {
+        "url": "gdrive://root/data",
+        "gdrive_client_id": "client",
+        "gdrive_client_secret": "secret",
+    }
+
+    def test_init(self):
+        remote = RemoteGDrive(Repo(), self.CONFIG)
+        assert str(remote.path_info) == self.CONFIG["url"]
+
+    def test_drive(self):
+        remote = RemoteGDrive(Repo(), self.CONFIG)
+        os.environ[
+            RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA
+        ] = USER_CREDS_TOKEN_REFRESH_ERROR
+        with pytest.raises(GDriveAccessTokenRefreshError):
+            remote.drive
+
+        os.environ[RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA] = ""
+        remote = RemoteGDrive(Repo(), self.CONFIG)
+        os.environ[
+            RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA
+        ] = USER_CREDS_MISSED_KEY_ERROR
+        with pytest.raises(GDriveMissedCredentialKeyError):
+            remote.drive