mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-11 07:18:28 +01:00
Merge pull request #2991 from MaxRis/gdrive-tests
Improve Google Drive auth failure error reporting
This commit is contained in:
+36
-7
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user