def _config_drive_folder_map_or_exit(config):
"""
Lists folders under our top level parent for this environment and returns
a dict of {partner name: folder id}. Partner names should match the values
in config['org_partner_mapping']
"""
drive = DriveApi(config['google_secrets_file'])
try:
LOG('Attempting to find all partner sub-directories on Drive.')
folders = drive.walk_files(
config['drive_partners_folder'],
mimetype='application/vnd.google-apps.folder',
recurse=False)
except Exception as exc: # pylint: disable=broad-except
FAIL_EXCEPTION(ERR_DRIVE_LISTING,
'Finding partner directories on Drive failed.', exc)
if not folders:
FAIL(
ERR_DRIVE_LISTING,
'Finding partner directories on Drive failed. Check your permissions.'
)
# As in _config_or_exit we force normalize the unicode here to make sure the keys
# match. Otherwise the name we get back from Google won't match what's in the YAML config.
config['partner_folder_mapping'] = OrderedDict()
for folder in folders:
folder['name'] = unicodedata.normalize('NFKC',
text_type(folder['name']))
config['partner_folder_mapping'][folder['name']] = folder['id']
def test_walk_files_retry(self, mock_from_service_account_file): # pylint: disable=unused-argument
"""
Subfolders are requested, but there is rate limiting causing a retry.
"""
fake_folders = [
{
'id': 'fake-folder-id-{}'.format(idx),
'name': 'fake-folder-name-{}'.format(idx),
'mimeType': 'application/vnd.google-apps.folder'
}
for idx in range(10)
]
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 list files, but the response suggests to retry.
self._http_mock_sequence_retry(),
# Finally, the request is retried, and the response is OK.
(
{'status': '200', 'content-type': 'application/json'},
json.dumps({'files': fake_folders}).encode('utf-8'),
),
])
test_client = DriveApi('non-existent-secrets.json', http=http_mock_sequence)
response = test_client.walk_files('fake-folder-id', mimetype=FOLDER_MIMETYPE, recurse=False)
# Remove all the mimeTypes for comparison purposes.
for fake_folder in fake_folders:
del fake_folder['mimeType']
six.assertCountEqual(self, response, fake_folders)
def test_walk_files_multi_page_csv_only(self, mock_from_service_account_file): # pylint: disable=unused-argument
"""
Files are searched for - and returned in two pages.
"""
fake_folder = [
{
'id': 'fake-folder-id-0',
'name': 'fake-folder-name-0',
'mimeType': 'application/vnd.google-apps.folder'
}
]
fake_csv_files = [
{
'id': 'fake-csv-file-id-{}'.format(idx),
'name': 'fake-csv-file-name-{}'.format(idx),
'mimeType': 'application/csv'
}
for idx in range(10)
]
fake_files_part_1 = fake_folder + fake_csv_files[:3]
fake_files_part_2 = fake_csv_files[3:8]
fake_files_part_3 = fake_csv_files[8:]
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 list files. The response contains a single folder and other files.
(
{'status': '200', 'content-type': 'application/json'},
json.dumps({'files': fake_files_part_1}).encode('utf-8'),
),
# Then, a request is made to list files from the single found folder.
# The response contains a nextPageToken indicating there are more pages.
(
{'status': '200', 'content-type': 'application/json'},
json.dumps({'files': fake_files_part_2, 'nextPageToken': 'fake-next-page-token'}).encode('utf-8'),
),
# Finally, another list request is made. This time, no nextPageToken is present in the response,
# indicating there are no more pages.
(
{'status': '200', 'content-type': 'application/json'},
json.dumps({'files': fake_files_part_3}).encode('utf-8'),
),
])
test_client = DriveApi('non-existent-secrets.json', http=http_mock_sequence)
response = test_client.walk_files('fake-folder-id', mimetype='application/csv')
# Remove all the mimeTypes for comparison purposes.
for fake_file in fake_csv_files:
del fake_file['mimeType']
six.assertCountEqual(self, response, fake_csv_files)
def test_walk_files_two_page(self, mock_from_service_account_file): # pylint: disable=unused-argument
"""
Subfolders are requested, but the response is paginated.
"""
fake_folders = [
{
'id': 'fake-folder-id-{}'.format(idx),
'name': 'fake-folder-name-{}'.format(idx),
'mimeType': 'application/vnd.google-apps.folder'
}
for idx in range(10)
]
fake_files_part_1 = fake_folders[:7]
fake_files_part_2 = fake_folders[7:]
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 list files. The response contains a nextPageToken suggesting there are more
# pages.
(
{'status': '200', 'content-type': 'application/json'},
json.dumps({'files': fake_files_part_1, 'nextPageToken': 'fake-next-page-token'}).encode('utf-8'),
),
# Finally, a second list request is made. This time, no nextPageToken is present in the response,
# suggesting there are no more pages.
(
{'status': '200', 'content-type': 'application/json'},
json.dumps({'files': fake_files_part_2}).encode('utf-8'),
),
])
test_client = DriveApi('non-existent-secrets.json', http=http_mock_sequence)
response = test_client.walk_files('fake-folder-id', mimetype=FOLDER_MIMETYPE, recurse=False)
# Remove all the mimeTypes for comparison purposes.
for fake_folder in fake_folders:
del fake_folder['mimeType']
six.assertCountEqual(self, response, fake_folders)