def test_retrieve_mail_attachments_with_regex_found(self, mock_imaplib): _create_fake_imap(mock_imaplib, with_mail=True) with ImapHook() as imap_hook: attachments_in_inbox = imap_hook.retrieve_mail_attachments( name=r'test(\d+).csv', check_regex=True) self.assertEqual(attachments_in_inbox, [('test1.csv', b'SWQsTmFtZQoxLEZlbGl4')])
def test_connect_and_disconnect(self, mock_imaplib): mock_conn = _create_fake_imap(mock_imaplib) with ImapHook(): pass mock_imaplib.IMAP4_SSL.assert_called_once_with('imap_server_address') mock_conn.login.assert_called_once_with('imap_user', 'imap_password') # pylint: disable=no-member assert mock_conn.logout.call_count == 1
def test_download_mail_attachments_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, 'test1.txt', 'test_directory') mock_open_method.assert_not_called() mock_open_method.return_value.write.assert_not_called()
def test_retrieve_mail_attachments_latest_only(self, mock_imaplib): _create_fake_imap(mock_imaplib, with_mail=True) with ImapHook() as imap_hook: attachments_in_inbox = imap_hook.retrieve_mail_attachments( name='test1.csv', latest_only=True) self.assertEqual(attachments_in_inbox, [('test1.csv', b'SWQsTmFtZQoxLEZlbGl4')])
def test_retrieve_mail_attachments_with_regex_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, name=r'test_(\d+).csv', check_regex=True)
def test_has_mail_attachment_with_regex_found(self, mock_imaplib): _create_fake_imap(mock_imaplib, with_mail=True) with ImapHook() as imap_hook: has_attachment_in_inbox = imap_hook.has_mail_attachment( name=r'test(\d+).csv', check_regex=True ) self.assertTrue(has_attachment_in_inbox)
def test_retrieve_mail_attachments_with_mail_filter(self, mock_imaplib): _create_fake_imap(mock_imaplib, with_mail=True) mail_filter = '(SINCE "01-Jan-2019")' with ImapHook() as imap_hook: imap_hook.retrieve_mail_attachments(name='test1.csv', mail_filter=mail_filter) mock_imaplib.IMAP4_SSL.return_value.search.assert_called_once_with( None, mail_filter)
def test_download_mail_attachments_with_escaping_chars(self, mock_imaplib, mock_open_method): _create_fake_imap(mock_imaplib, with_mail=True, attachment_name='../test1.csv') with ImapHook() as imap_hook: imap_hook.download_mail_attachments( name='../test1.csv', local_output_directory='test_directory' ) mock_open_method.assert_not_called() mock_open_method.return_value.write.assert_not_called()
def test_download_mail_attachments_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('test1.csv', 'test_directory') 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_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_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 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 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')