def test_comment_files_with_nonexistent_file(self, mock_from_service_account_file): # pylint: disable=unused-argument
"""
Test case for commenting on files, where some files are nonexistent.
"""
fake_file_ids = ['fake-file-id0', 'fake-file-id1']
batch_response = b'''--batch_foobarbaz
Content-Type: application/http
Content-Transfer-Encoding: binary
Content-ID: <response+0>
HTTP/1.1 404 NOT FOUND
Content-Type: application/json
Content-length: 266
ETag: "etag/pony"\r\n\r\n{
"error": {
"errors": [
{
"domain": "global",
"reason": "notFound",
"message": "File not found: fake-file-id0.",
"locationType": "parameter",
"location": "fileId"
}
],
"code": 404,
"message": "File not found: fake-file-id0."
}
}
--batch_foobarbaz
Content-Type: application/http
Content-Transfer-Encoding: binary
Content-ID: <response+1>
HTTP/1.1 204 OK
ETag: "etag/sheep"\r\n\r\n{"id": "fake-comment-id1"}
--batch_foobarbaz--'''
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 add comments to the files.
({'status': '200', 'content-type': 'multipart/mixed; boundary="batch_foobarbaz"'}, batch_response),
])
test_client = DriveApi('non-existent-secrets.json', http=http_mock_sequence)
if sys.version_info < (3, 4):
# This is a simple smoke-test without checking the output because python <3.4 doesn't support assertLogs.
with self.assertRaises(BatchRequestError):
test_client.create_comments_for_files(list(zip(fake_file_ids, cycle(['some comment message']))))
else:
# This is the full test case, which only runs under python 3.4+.
with self.assertLogs(level='INFO') as captured_logs: # pylint: disable=no-member
with self.assertRaises(BatchRequestError):
test_client.create_comments_for_files(list(zip(fake_file_ids, cycle(['some comment message']))))
assert sum('Successfully processed request' in msg for msg in captured_logs.output) == 1
assert sum('Error processing request' in msg for msg in captured_logs.output) == 1
def test_comment_files_with_duplicate_file(self, mock_from_service_account_file): # pylint: disable=unused-argument
"""
Test case for duplicate file IDs.
"""
fake_file_ids = ['fake-file-id0', 'fake-file-id1', 'fake-file-id0']
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),
])
test_client = DriveApi('non-existent-secrets.json', http=http_mock_sequence)
with self.assertRaises(ValueError):
test_client.create_comments_for_files(list(zip(fake_file_ids, cycle(['some comment message']))))
def _add_comments_to_files(config, file_ids):
"""
Add comments to the uploaded csv files, triggering email notification.
Args:
file_ids (dict): Mapping of partner names to Drive file IDs corresponding to the newly uploaded csv files.
"""
drive = DriveApi(config['google_secrets_file'])
partner_folders_to_permissions = drive.list_permissions_for_files(
config['partner_folder_mapping'].values(),
fields='emailAddress',
)
# create a mapping of partners to a list of permissions dicts:
permissions = {
partner: partner_folders_to_permissions[
config['partner_folder_mapping'][partner]]
for partner in file_ids
}
# throw out all blacklisted addresses, and flatten the permissions dicts to just the email:
external_emails = {
partner: [
perm['emailAddress'] for perm in permissions[partner]
if not any(perm['emailAddress'].lower().endswith(
blacklisted_domain.lower()) for blacklisted_domain in
config['blacklisted_notification_domains'])
]
for partner in permissions
}
file_ids_and_comments = []
for partner in file_ids:
if not external_emails[partner]:
LOG('WARNING: could not find a POC for the following partner: "{}". '
'Double check the partner folder permissions in Google Drive.'.
format(partner))
else:
tag_string = ' '.join('+' + email
for email in external_emails[partner])
comment_content = NOTIFICATION_MESSAGE_TEMPLATE.format(
tags=tag_string)
file_ids_and_comments.append((file_ids[partner], comment_content))
try:
LOG('Adding notification comments to uploaded csv files.')
drive.create_comments_for_files(file_ids_and_comments)
except Exception as exc: # pylint: disable=broad-except
# do not fail the script here, since comment errors are non-critical
LOG('WARNING: there was an error adding Google Drive comments to the csv files: {}'
.format(exc))
def test_comment_files_success(self, mock_from_service_account_file): # pylint: disable=unused-argument
"""
Test normal case for commenting on files.
"""
fake_file_ids = ['fake-file-id0', 'fake-file-id1']
batch_response = b'''--batch_foobarbaz
Content-Type: application/http
Content-Transfer-Encoding: binary
Content-ID: <response+0>
HTTP/1.1 204 OK
ETag: "etag/pony"\r\n\r\n{"id": "fake-comment-id0"}
--batch_foobarbaz
Content-Type: application/http
Content-Transfer-Encoding: binary
Content-ID: <response+1>
HTTP/1.1 204 OK
ETag: "etag/sheep"\r\n\r\n{"id": "fake-comment-id1"}
--batch_foobarbaz--'''
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 add comments to the files.
({'status': '200', 'content-type': 'multipart/mixed; boundary="batch_foobarbaz"'}, batch_response),
])
test_client = DriveApi('non-existent-secrets.json', http=http_mock_sequence)
resp = test_client.create_comments_for_files(list(zip(fake_file_ids, cycle(['some comment message']))))
six.assertCountEqual(
self,
resp,
{
'fake-file-id0': {'id': 'fake-comment-id0'},
'fake-file-id1': {'id': 'fake-comment-id1'},
},
)
def test_comment_files_batching_retries(self, mock_from_service_account_file): # pylint: disable=unused-argument
"""
Test commenting on more files than the google API batch limit. This also tests the partial retry
mechanism when a subset of responses are rate limited.
"""
num_files = int(GOOGLE_API_MAX_BATCH_SIZE * 1.5)
fake_file_ids = ['fake-file-id{}'.format(n) for n in range(num_files)]
batch_response_0 = '\n'.join(
'''--batch_foobarbaz
Content-Type: application/http
Content-Transfer-Encoding: binary
Content-ID: <response+{idx}>
HTTP/1.1 204 OK
ETag: "etag/pony{idx}"\r\n\r\n{{"id": "fake-comment-id{idx}"}}
'''.format(idx=n)
for n in range(GOOGLE_API_MAX_BATCH_SIZE)
)
batch_response_0 += '--batch_foobarbaz--'
batch_response_1 = '\n'.join(
'''--batch_foobarbaz
Content-Type: application/http
Content-Transfer-Encoding: binary
Content-ID: <response+{idx}>
HTTP/1.1 204 OK
ETag: "etag/pony{idx}"\r\n\r\n{{"id": "fake-comment-id{idx}"}}
'''.format(idx=n)
for n in range(int(GOOGLE_API_MAX_BATCH_SIZE * 0.25))
)
batch_response_1 += '\n'
batch_response_1 += '\n'.join(
'''--batch_foobarbaz
Content-Type: application/http
Content-Transfer-Encoding: binary
Content-ID: <response+{idx}>
HTTP/1.1 500 Internal Server Error
ETag: "etag/pony{idx}"\r\n\r\n
'''.format(idx=n)
for n in range(int(GOOGLE_API_MAX_BATCH_SIZE * 0.25), int(GOOGLE_API_MAX_BATCH_SIZE * 0.5))
)
batch_response_1 += '--batch_foobarbaz--'
batch_response_2 = '\n'.join(
'''--batch_foobarbaz
Content-Type: application/http
Content-Transfer-Encoding: binary
Content-ID: <response+{idx}>
HTTP/1.1 204 OK
ETag: "etag/pony{idx}"\r\n\r\n{{"id": "fake-comment-id{idx}"}}
'''.format(idx=n)
for n in range(int(GOOGLE_API_MAX_BATCH_SIZE * 0.25), int(GOOGLE_API_MAX_BATCH_SIZE * 0.5))
)
batch_response_2 += '--batch_foobarbaz--'
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 add comments to the files, first batch. Return max batch size results.
({'status': '200', 'content-type': 'multipart/mixed; boundary="batch_foobarbaz"'}, batch_response_0),
# Then, a request is made to add comments to the files, second batch. Only half of the results are returned,
# the rest resulted in HTTP 500.
({'status': '200', 'content-type': 'multipart/mixed; boundary="batch_foobarbaz"'}, batch_response_1),
# Then, a request is made retry the last half of the second batch (only the ones that returned the 500s).
# Return the last 1/4 results.
({'status': '200', 'content-type': 'multipart/mixed; boundary="batch_foobarbaz"'}, batch_response_2),
])
test_client = DriveApi('non-existent-secrets.json', http=http_mock_sequence)
resp = test_client.create_comments_for_files(list(zip(fake_file_ids, cycle(['some comment message']))))
six.assertCountEqual(
self,
resp,
{
'fake-file-id{}'.format(n): {'id': 'fake-comment-id{}'.format(n)}
for n in range(num_files)
},
)
