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