Repository: incubator-airflow
Updated Branches:
  refs/heads/master 3c4f1fd9e -> b3dea2f0d


[AIRFLOW-2281] Add support for Sendgrid categories

Closes #3188 from ms32035/sendgrid_categories


Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/b3dea2f0
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/b3dea2f0
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/b3dea2f0

Branch: refs/heads/master
Commit: b3dea2f0d4873bfce138131f9acd891868b32935
Parents: 3c4f1fd
Author: Marcin Szymanski <ms32...@gmail.com>
Authored: Mon Apr 9 10:51:25 2018 +0200
Committer: Fokko Driesprong <fokkodriespr...@godatadriven.com>
Committed: Mon Apr 9 10:51:25 2018 +0200

----------------------------------------------------------------------
 airflow/contrib/utils/sendgrid.py    | 14 ++++++++++----
 tests/contrib/utils/test_sendgrid.py | 13 ++++++++-----
 2 files changed, 18 insertions(+), 9 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/b3dea2f0/airflow/contrib/utils/sendgrid.py
----------------------------------------------------------------------
diff --git a/airflow/contrib/utils/sendgrid.py 
b/airflow/contrib/utils/sendgrid.py
index 07614a5..ebc32ec 100644
--- a/airflow/contrib/utils/sendgrid.py
+++ b/airflow/contrib/utils/sendgrid.py
@@ -20,12 +20,13 @@ from __future__ import unicode_literals
 import base64
 import mimetypes
 import os
+
 import sendgrid
+from sendgrid.helpers.mail import Attachment, Content, Email, Mail, \
+    Personalization, CustomArg, Category
 
 from airflow.utils.email import get_email_address_list
 from airflow.utils.log.logging_mixin import LoggingMixin
-from sendgrid.helpers.mail import Attachment, Content, Email, Mail, \
-    Personalization, CustomArg
 
 
 def send_email(to, subject, html_content, files=None,
@@ -61,8 +62,6 @@ def send_email(to, subject, html_content, files=None,
         bcc = get_email_address_list(bcc)
         for bcc_address in bcc:
             personalization.add_bcc(Email(bcc_address))
-    mail.add_personalization(personalization)
-    mail.add_content(Content('text/html', html_content))
 
     # Add custom_args to personalization if present
     pers_custom_args = kwargs.get('personalization_custom_args', None)
@@ -70,6 +69,13 @@ def send_email(to, subject, html_content, files=None,
         for key in pers_custom_args.keys():
             personalization.add_custom_arg(CustomArg(key, 
pers_custom_args[key]))
 
+    mail.add_personalization(personalization)
+    mail.add_content(Content('text/html', html_content))
+
+    categories = kwargs.get('categories', [])
+    for cat in categories:
+        mail.add_category(Category(cat))
+
     # Add email attachment.
     for fname in files or []:
         basename = os.path.basename(fname)

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/b3dea2f0/tests/contrib/utils/test_sendgrid.py
----------------------------------------------------------------------
diff --git a/tests/contrib/utils/test_sendgrid.py 
b/tests/contrib/utils/test_sendgrid.py
index ba039e7..204aaa3 100644
--- a/tests/contrib/utils/test_sendgrid.py
+++ b/tests/contrib/utils/test_sendgrid.py
@@ -44,9 +44,11 @@ class SendEmailSendGridTest(unittest.TestCase):
             'from': {'email': u'f...@bar.com'},
             'subject': 'sendgrid-send-email unit test'}
         self.personalization_custom_args = {'arg1': 'val1', 'arg2': 'val2'}
-        self.expected_mail_data_custom_args = 
copy.deepcopy(self.expected_mail_data)
-        
self.expected_mail_data_custom_args['personalizations'][0]['custom_args'] = \
+        self.categories = ['cat1', 'cat2']
+        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
 
         # Test the right email is constructed.
 
@@ -60,8 +62,9 @@ class SendEmailSendGridTest(unittest.TestCase):
     # 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_custom_args(self, mock_post, 
mock_get):
+    def test_send_email_sendgrid_correct_email_extras(self, mock_post, 
mock_get):
         mock_get.return_value = 'f...@bar.com'
         send_email(self.to, self.subject, self.html_content, cc=self.cc, 
bcc=self.bcc,
-                   
personalization_custom_args=self.personalization_custom_args)
-        mock_post.assert_called_with(self.expected_mail_data_custom_args)
+                   
personalization_custom_args=self.personalization_custom_args,
+                   categories=self.categories)
+        mock_post.assert_called_with(self.expected_mail_data_extras)

Reply via email to