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 test_download_mail_attachments_with_latest_only( 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='test1.csv', local_output_directory='test_directory', latest_only=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_regex_not_found( self, mock_imaplib, mock_open_method): _create_fake_imap(mock_imaplib, with_mail=True) with ImapHook() as imap_hook: self.assertRaises( AirflowException, imap_hook.download_mail_attachments, name=r'test_(\d+).csv', local_output_directory='test_directory', check_regex=True, ) 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, check_regex=self.check_regex, mail_folder=self.mail_folder, mail_filter=self.mail_filter)
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, check_regex=self.imap_check_regex, latest_only=True, mail_folder=self.imap_mail_folder, mail_filter=self.imap_mail_filter, ) s3_hook = S3Hook(aws_conn_id=self.s3_conn_id) s3_hook.load_bytes(bytes_data=imap_mail_attachments[0][1], key=self.s3_key, replace=self.s3_overwrite)
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')