def _push_files_to_google(config, partner_filenames):
"""
Copy the file to Google drive for this partner
Returns:
List of file IDs for the uploaded csv files.
"""
# First make sure we have Drive folders for all partners
failed_partners = []
for partner in partner_filenames:
if partner not in config['partner_folder_mapping']:
failed_partners.append(partner)
if failed_partners:
FAIL(ERR_BAD_CONFIG, 'These partners have retiring learners, but no Drive folder: {}'.format(failed_partners))
file_ids = {}
drive = DriveApi(config['google_secrets_file'])
for partner in partner_filenames:
# This is populated on the fly in _config_drive_folder_map_or_exit
folder_id = config['partner_folder_mapping'][partner]
file_id = None
with open(partner_filenames[partner], 'rb') as f:
try:
drive_filename = os.path.basename(partner_filenames[partner])
LOG('Attempting to upload {} to {} Drive folder.'.format(drive_filename, partner))
file_id = drive.create_file_in_folder(folder_id, drive_filename, f, "text/csv")
except Exception as exc: # pylint: disable=broad-except
FAIL_EXCEPTION(ERR_DRIVE_UPLOAD, 'Drive upload failed for: {}'.format(drive_filename), exc)
file_ids[partner] = file_id
return file_ids
def test_create_file_success(self, mock_from_service_account_file): # pylint: disable=unused-argument
"""
Test normal case for uploading a file.
"""
fake_file_id = 'fake-file-id'
http_mock_sequence = HttpMockSequence([
# First, a request is made to the discovery API to construct a client object for Drive.
({'status': '200'}, self.mock_discovery_response_content),
# Then, a request is made to upload the file.
({'status': '200'}, '{{"id": "{}"}}'.format(fake_file_id)),
])
test_client = DriveApi('non-existent-secrets.json', http=http_mock_sequence)
response = test_client.create_file_in_folder(
'fake-folder-id',
'Fake Filename',
BytesIO('fake file contents'.encode('ascii')),
'text/plain',
)
assert response == fake_file_id
def test_create_file_retry_success(self, mock_from_service_account_file):
"""
Test rate limit and retry during file upload.
"""
fake_file_id = 'fake-file-id'
http_mock_sequence = HttpMockSequence([
# First, a request is made to the discovery API to construct a client object for Drive.
({'status': '200'}, self.mock_discovery_response_content),
# Then, a request is made to upload the file while rate limiting was activated. This should cause a retry.
self._http_mock_sequence_retry(),
# Finally, success.
({'status': '200'}, '{{"id": "{}"}}'.format(fake_file_id)),
])
test_client = DriveApi('non-existent-secrets.json', http=http_mock_sequence)
response = test_client.create_file_in_folder(
'fake-folder-id',
'Fake Filename',
BytesIO('fake file contents'.encode('ascii')),
'text/plain',
)
# There is no need to explicitly check if the call was retried because
# the response value cannot possibly contain fake_file_id otherwise,
# since it was only passed in the last response.
assert response == fake_file_id