Merge pull request #2991 from MaxRis/gdrive-tests

Improve Google Drive auth failure error reporting
This commit is contained in:
Ruslan Kuprieiev
2019-12-29 04:35:29 +02:00
committed by GitHub
2 changed files with 80 additions and 14 deletions
+36 -7
View File
@@ -21,8 +21,18 @@ FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"
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 @@ class RemoteGDrive(RemoteBASE):
"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 @@ class RemoteGDrive(RemoteBASE):
@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 @@ class RemoteGDrive(RemoteBASE):
# 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)
+44 -7
View File
@@ -1,9 +1,46 @@
import mock
from dvc.remote.gdrive import RemoteGDrive
from unittest import TestCase
import pytest
import os
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 = "{}"
@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
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