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