Repository: incubator-airflow Updated Branches: refs/heads/master e44688ed0 -> b0d0d0a04
[AIRFLOW-2377] Improve Sendgrid sender support Closes #3266 from ms32035/sendgrid_sender Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/b0d0d0a0 Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/b0d0d0a0 Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/b0d0d0a0 Branch: refs/heads/master Commit: b0d0d0a041bd5c4b19f6b97dc0fa289436743272 Parents: e44688e Author: Marcin Szymanski <[email protected]> Authored: Thu Apr 26 10:25:32 2018 -0700 Committer: r39132 <[email protected]> Committed: Thu Apr 26 10:25:32 2018 -0700 ---------------------------------------------------------------------- airflow/contrib/utils/sendgrid.py | 8 +++++--- tests/contrib/utils/test_sendgrid.py | 31 +++++++++++++++++++++++++++---- 2 files changed, 32 insertions(+), 7 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/b0d0d0a0/airflow/contrib/utils/sendgrid.py ---------------------------------------------------------------------- diff --git a/airflow/contrib/utils/sendgrid.py b/airflow/contrib/utils/sendgrid.py index 206bb93..9055c97 100644 --- a/airflow/contrib/utils/sendgrid.py +++ b/airflow/contrib/utils/sendgrid.py @@ -7,9 +7,9 @@ # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at -# +# # http://www.apache.org/licenses/LICENSE-2.0 -# +# # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY @@ -51,7 +51,9 @@ def send_email(to, subject, html_content, files=None, SENDGRID_API_KEY={your-sendgrid-api-key}. """ mail = Mail() - mail.from_email = Email(os.environ.get('SENDGRID_MAIL_FROM')) + from_email = kwargs.get('from_email') or os.environ.get('SENDGRID_MAIL_FROM') + from_name = kwargs.get('from_name') or os.environ.get('SENDGRID_MAIL_SENDER') + mail.from_email = Email(from_email, from_name) mail.subject = subject # Add the recipient list of to emails. http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/b0d0d0a0/tests/contrib/utils/test_sendgrid.py ---------------------------------------------------------------------- diff --git a/tests/contrib/utils/test_sendgrid.py b/tests/contrib/utils/test_sendgrid.py index 997bc55..6710076 100644 --- a/tests/contrib/utils/test_sendgrid.py +++ b/tests/contrib/utils/test_sendgrid.py @@ -7,9 +7,9 @@ # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at -# +# # http://www.apache.org/licenses/LICENSE-2.0 -# +# # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY @@ -50,17 +50,27 @@ class SendEmailSendGridTest(unittest.TestCase): 'subject': 'sendgrid-send-email unit test'} self.personalization_custom_args = {'arg1': 'val1', 'arg2': 'val2'} self.categories = ['cat1', 'cat2'] + # extras self.expected_mail_data_extras = copy.deepcopy(self.expected_mail_data) self.expected_mail_data_extras['personalizations'][0]['custom_args'] = \ self.personalization_custom_args self.expected_mail_data_extras['categories'] = self.categories + self.expected_mail_data_extras['from'] = \ + {'name': 'Foo', 'email': '[email protected]'} + # sender + self.expected_mail_data_sender = copy.deepcopy(self.expected_mail_data) + self.expected_mail_data_sender['from'] = \ + {'name': 'Foo Bar', 'email': '[email protected]'} # Test the right email is constructed. @mock.patch('os.environ.get') @mock.patch('airflow.contrib.utils.sendgrid._post_sendgrid_mail') def test_send_email_sendgrid_correct_email(self, mock_post, mock_get): - mock_get.return_value = '[email protected]' + def get_return(var): + return {'SENDGRID_MAIL_FROM': '[email protected]'}.get(var) + + mock_get.side_effect = get_return send_email(self.to, self.subject, self.html_content, cc=self.cc, bcc=self.bcc) mock_post.assert_called_with(self.expected_mail_data) @@ -68,8 +78,21 @@ class SendEmailSendGridTest(unittest.TestCase): @mock.patch('os.environ.get') @mock.patch('airflow.contrib.utils.sendgrid._post_sendgrid_mail') def test_send_email_sendgrid_correct_email_extras(self, mock_post, mock_get): - mock_get.return_value = '[email protected]' + def get_return(var): + return {'SENDGRID_MAIL_FROM': '[email protected]', + 'SENDGRID_MAIL_SENDER': 'Foo'}.get(var) + + mock_get.side_effect = get_return send_email(self.to, self.subject, self.html_content, cc=self.cc, bcc=self.bcc, personalization_custom_args=self.personalization_custom_args, categories=self.categories) mock_post.assert_called_with(self.expected_mail_data_extras) + + @mock.patch('os.environ.get') + @mock.patch('airflow.contrib.utils.sendgrid._post_sendgrid_mail') + def test_send_email_sendgrid_sender(self, mock_post, mock_get): + + mock_get.return_value = None + send_email(self.to, self.subject, self.html_content, cc=self.cc, bcc=self.bcc, + from_email='[email protected]', from_name='Foo Bar') + mock_post.assert_called_with(self.expected_mail_data_sender)
