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_list_permissions_success(self, mock_from_service_account_file): # pylint: disable=unused-argument
"""
Test normal case for listing permissions 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 200 OK
Content-Type: application/json
ETag: "etag/pony"\r\n\r\n{"permissions": [{"emailAddress": "*****@*****.**", "role": "reader"}]}
--batch_foobarbaz
Content-Type: application/http
Content-Transfer-Encoding: binary
Content-ID: <response+1>
HTTP/1.1 200 OK
Content-Type: application/json
ETag: "etag/sheep"\r\n\r\n{"permissions": [{"emailAddress": "*****@*****.**", "role": "writer"}]}
--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.list_permissions_for_files(fake_file_ids)
six.assertCountEqual(
self,
resp,
{
'fake-file-id0': [{'emailAddress': '*****@*****.**', 'role': 'reader'}],
'fake-file-id1': [{'emailAddress': '*****@*****.**', 'role': 'writer'}],
},
)
def test_list_permissions_one_failure(self, mock_from_service_account_file): # pylint: disable=unused-argument
"""
Test case for listing permissions on files, but one file doesn't exist.
"""
fake_file_ids = ['fake-file-id0', 'fake-file-id1', 'fake-file-id2']
batch_response = b'''--batch_foobarbaz
Content-Type: application/http
Content-Transfer-Encoding: binary
Content-ID: <response+0>
HTTP/1.1 200 OK
Content-Type: application/json
ETag: "etag/pony"\r\n\r\n{"permissions": [{"emailAddress": "*****@*****.**", "role": "reader"}]}
--batch_foobarbaz
Content-Type: application/http
Content-Transfer-Encoding: binary
Content-ID: <response+1>
HTTP/1.1 200 OK
Content-Type: application/json
ETag: "etag/sheep"\r\n\r\n{"permissions": [{"emailAddress": "*****@*****.**", "role": "writer"}]}
--batch_foobarbaz
Content-Type: application/http
Content-Transfer-Encoding: binary
Content-ID: <response+2>
HTTP/1.1 404 NOT FOUND
Content-Type: application/json
Content-length: 266
ETag: "etag/bird"\r\n\r\n{
"error": {
"errors": [
{
"domain": "global",
"reason": "notFound",
"message": "File not found: fake-file-id2.",
"locationType": "parameter",
"location": "fileId"
}
],
"code": 404,
"message": "File not found: fake-file-id2."
}
}
--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.list_permissions_for_files(fake_file_ids)
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.list_permissions_for_files(fake_file_ids)
assert sum('Successfully processed request' in msg for msg in captured_logs.output) == 2
assert sum('Error processing request' in msg for msg in captured_logs.output) == 1