forked from apache/airflow
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[AIRFLOW-5896] Move email stuff from tests/core.py (apache#6545)
* [AIRFLOW-5896] Move email stuff from tests/ccore.py
- Loading branch information
Showing
2 changed files
with
137 additions
and
135 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -16,13 +16,23 @@ | |
| # KIND, either express or implied. See the License for the | ||
| # specific language governing permissions and limitations | ||
| # under the License. | ||
|
|
||
| import os | ||
| import tempfile | ||
| import unittest | ||
| from email.mime.application import MIMEApplication | ||
| from email.mime.multipart import MIMEMultipart | ||
| from email.mime.text import MIMEText | ||
|
|
||
| import mock | ||
|
|
||
| from airflow import conf, utils | ||
| from airflow.utils.email import get_email_address_list | ||
| from tests import conf_vars | ||
|
|
||
| EMAILS = ['[email protected]', '[email protected]'] | ||
|
|
||
| send_email_test = mock.MagicMock() | ||
|
|
||
|
|
||
| class TestEmail(unittest.TestCase): | ||
|
|
||
|
|
@@ -67,3 +77,129 @@ def test_get_email_address_invalid_type_in_iterable(self): | |
|
|
||
| self.assertRaises( | ||
| TypeError, get_email_address_list, emails_list) | ||
|
|
||
| def setUp(self): | ||
| conf.remove_option('email', 'EMAIL_BACKEND') | ||
|
|
||
| @mock.patch('airflow.utils.email.send_email') | ||
| def test_default_backend(self, mock_send_email): | ||
| res = utils.email.send_email('to', 'subject', 'content') | ||
| mock_send_email.assert_called_once_with('to', 'subject', 'content') | ||
| self.assertEqual(mock_send_email.return_value, res) | ||
|
|
||
| @mock.patch('airflow.utils.email.send_email_smtp') | ||
| def test_custom_backend(self, mock_send_email): | ||
| with conf_vars({('email', 'email_backend'): 'tests.utils.test_email.send_email_test'}): | ||
| utils.email.send_email('to', 'subject', 'content') | ||
| send_email_test.assert_called_once_with( | ||
| 'to', 'subject', 'content', files=None, dryrun=False, | ||
| cc=None, bcc=None, mime_charset='utf-8', mime_subtype='mixed') | ||
| self.assertFalse(mock_send_email.called) | ||
|
|
||
|
|
||
| class TestEmailSmtp(unittest.TestCase): | ||
| def setUp(self): | ||
| conf.set('smtp', 'SMTP_SSL', 'False') | ||
|
|
||
| @mock.patch('airflow.utils.email.send_MIME_email') | ||
| def test_send_smtp(self, mock_send_mime): | ||
| attachment = tempfile.NamedTemporaryFile() | ||
| attachment.write(b'attachment') | ||
| attachment.seek(0) | ||
| utils.email.send_email_smtp('to', 'subject', 'content', files=[attachment.name]) | ||
| self.assertTrue(mock_send_mime.called) | ||
| call_args = mock_send_mime.call_args[0] | ||
| self.assertEqual(conf.get('smtp', 'SMTP_MAIL_FROM'), call_args[0]) | ||
| self.assertEqual(['to'], call_args[1]) | ||
| msg = call_args[2] | ||
| self.assertEqual('subject', msg['Subject']) | ||
| self.assertEqual(conf.get('smtp', 'SMTP_MAIL_FROM'), msg['From']) | ||
| self.assertEqual(2, len(msg.get_payload())) | ||
| filename = 'attachment; filename="' + os.path.basename(attachment.name) + '"' | ||
| self.assertEqual(filename, msg.get_payload()[-1].get('Content-Disposition')) | ||
| mimeapp = MIMEApplication('attachment') | ||
| self.assertEqual(mimeapp.get_payload(), msg.get_payload()[-1].get_payload()) | ||
|
|
||
| @mock.patch('airflow.utils.email.send_MIME_email') | ||
| def test_send_smtp_with_multibyte_content(self, mock_send_mime): | ||
| utils.email.send_email_smtp('to', 'subject', '🔥', mime_charset='utf-8') | ||
| self.assertTrue(mock_send_mime.called) | ||
| call_args = mock_send_mime.call_args[0] | ||
| msg = call_args[2] | ||
| mimetext = MIMEText('🔥', 'mixed', 'utf-8') | ||
| self.assertEqual(mimetext.get_payload(), msg.get_payload()[0].get_payload()) | ||
|
|
||
| @mock.patch('airflow.utils.email.send_MIME_email') | ||
| def test_send_bcc_smtp(self, mock_send_mime): | ||
| attachment = tempfile.NamedTemporaryFile() | ||
| attachment.write(b'attachment') | ||
| attachment.seek(0) | ||
| utils.email.send_email_smtp('to', 'subject', 'content', files=[attachment.name], cc='cc', bcc='bcc') | ||
| self.assertTrue(mock_send_mime.called) | ||
| call_args = mock_send_mime.call_args[0] | ||
| self.assertEqual(conf.get('smtp', 'SMTP_MAIL_FROM'), call_args[0]) | ||
| self.assertEqual(['to', 'cc', 'bcc'], call_args[1]) | ||
| msg = call_args[2] | ||
| self.assertEqual('subject', msg['Subject']) | ||
| self.assertEqual(conf.get('smtp', 'SMTP_MAIL_FROM'), msg['From']) | ||
| self.assertEqual(2, len(msg.get_payload())) | ||
| self.assertEqual('attachment; filename="' + os.path.basename(attachment.name) + '"', | ||
| msg.get_payload()[-1].get('Content-Disposition')) | ||
| mimeapp = MIMEApplication('attachment') | ||
| self.assertEqual(mimeapp.get_payload(), msg.get_payload()[-1].get_payload()) | ||
|
|
||
| @mock.patch('smtplib.SMTP_SSL') | ||
| @mock.patch('smtplib.SMTP') | ||
| def test_send_mime(self, mock_smtp, mock_smtp_ssl): | ||
| mock_smtp.return_value = mock.Mock() | ||
| mock_smtp_ssl.return_value = mock.Mock() | ||
| msg = MIMEMultipart() | ||
| utils.email.send_MIME_email('from', 'to', msg, dryrun=False) | ||
| mock_smtp.assert_called_once_with( | ||
| conf.get('smtp', 'SMTP_HOST'), | ||
| conf.getint('smtp', 'SMTP_PORT'), | ||
| ) | ||
| self.assertTrue(mock_smtp.return_value.starttls.called) | ||
| mock_smtp.return_value.login.assert_called_once_with( | ||
| conf.get('smtp', 'SMTP_USER'), | ||
| conf.get('smtp', 'SMTP_PASSWORD'), | ||
| ) | ||
| mock_smtp.return_value.sendmail.assert_called_once_with('from', 'to', msg.as_string()) | ||
| self.assertTrue(mock_smtp.return_value.quit.called) | ||
|
|
||
| @mock.patch('smtplib.SMTP_SSL') | ||
| @mock.patch('smtplib.SMTP') | ||
| def test_send_mime_ssl(self, mock_smtp, mock_smtp_ssl): | ||
| mock_smtp.return_value = mock.Mock() | ||
| mock_smtp_ssl.return_value = mock.Mock() | ||
| with conf_vars({('smtp', 'smtp_ssl'): 'True'}): | ||
| utils.email.send_MIME_email('from', 'to', MIMEMultipart(), dryrun=False) | ||
| self.assertFalse(mock_smtp.called) | ||
| mock_smtp_ssl.assert_called_once_with( | ||
| conf.get('smtp', 'SMTP_HOST'), | ||
| conf.getint('smtp', 'SMTP_PORT'), | ||
| ) | ||
|
|
||
| @mock.patch('smtplib.SMTP_SSL') | ||
| @mock.patch('smtplib.SMTP') | ||
| def test_send_mime_noauth(self, mock_smtp, mock_smtp_ssl): | ||
| mock_smtp.return_value = mock.Mock() | ||
| mock_smtp_ssl.return_value = mock.Mock() | ||
| with conf_vars({ | ||
| ('smtp', 'smtp_user'): None, | ||
| ('smtp', 'smtp_password'): None, | ||
| }): | ||
| utils.email.send_MIME_email('from', 'to', MIMEMultipart(), dryrun=False) | ||
| self.assertFalse(mock_smtp_ssl.called) | ||
| mock_smtp.assert_called_once_with( | ||
| conf.get('smtp', 'SMTP_HOST'), | ||
| conf.getint('smtp', 'SMTP_PORT'), | ||
| ) | ||
| self.assertFalse(mock_smtp.login.called) | ||
|
|
||
| @mock.patch('smtplib.SMTP_SSL') | ||
| @mock.patch('smtplib.SMTP') | ||
| def test_send_mime_dryrun(self, mock_smtp, mock_smtp_ssl): | ||
| utils.email.send_MIME_email('from', 'to', MIMEMultipart(), dryrun=True) | ||
| self.assertFalse(mock_smtp.called) | ||
| self.assertFalse(mock_smtp_ssl.called) | ||