def test_download_mail_attachments_with_regex_found(
self, mock_imaplib, mock_open_method):
_create_fake_imap(mock_imaplib, with_mail=True)
with ImapHook() as imap_hook:
imap_hook.download_mail_attachments(
name=r'test(\d+).csv',
local_output_directory='test_directory',
check_regex=True)
mock_open_method.assert_called_once_with('test_directory/test1.csv',
'wb')
mock_open_method.return_value.write.assert_called_once_with(
b'SWQsTmFtZQoxLEZlbGl4')
def test_download_mail_attachments_with_mail_filter(
self, mock_imaplib, mock_open_method):
_create_fake_imap(mock_imaplib, with_mail=True)
mail_filter = '(SINCE "01-Jan-2019")'
with ImapHook() as imap_hook:
imap_hook.download_mail_attachments(
name='test1.csv',
local_output_directory='test_directory',
mail_filter=mail_filter)
mock_imaplib.IMAP4_SSL.return_value.search.assert_called_once_with(
None, mail_filter)
assert mock_open_method.call_count == 1
def test_download_mail_attachments_with_symlink(self, mock_imaplib,
mock_open_method,
mock_is_symlink):
_create_fake_imap(mock_imaplib,
with_mail=True,
attachment_name='symlink')
with ImapHook() as imap_hook:
imap_hook.download_mail_attachments(
name='symlink', local_output_directory='test_directory')
assert mock_is_symlink.call_count == 1
mock_open_method.assert_not_called()
mock_open_method.return_value.write.assert_not_called()
def poke(self, context):
"""
Pokes for a mail attachment on the mail server.
:param context: The context that is being provided when poking.
:type context: dict
:return: True if attachment with the given name is present and False if not.
:rtype: bool
"""
self.log.info('Poking for %s', self.attachment_name)
with ImapHook(imap_conn_id=self.conn_id) as imap_hook:
return imap_hook.has_mail_attachment(name=self.attachment_name,
mail_folder=self.mail_folder,
check_regex=self.check_regex)
def execute(self, context):
"""
This function executes the transfer from the email server (via imap) into s3.
:param context: The context while executing.
:type context: dict
"""
self.log.info(
'Transferring mail attachment %s from mail server via imap to s3 key %s...',
self.imap_attachment_name, self.s3_key)
with ImapHook(imap_conn_id=self.imap_conn_id) as imap_hook:
imap_mail_attachments = imap_hook.retrieve_mail_attachments(
name=self.imap_attachment_name,
mail_folder=self.imap_mail_folder,
check_regex=self.imap_check_regex,
latest_only=True)
s3_hook = S3Hook(aws_conn_id=self.s3_conn_id)
s3_hook.load_string(string_data=imap_mail_attachments[0][1],
key=self.s3_key)
def test_retrieve_mail_attachments_not_found(self, mock_imaplib):
_create_fake_imap(mock_imaplib, with_mail=True)
with ImapHook() as imap_hook:
self.assertRaises(AirflowException, imap_hook.retrieve_mail_attachments, 'test1.txt')
