diff --git a/sendgrid/helpers/mail/mail.py b/sendgrid/helpers/mail/mail.py index ba21f789..50b79f73 100644 --- a/sendgrid/helpers/mail/mail.py +++ b/sendgrid/helpers/mail/mail.py @@ -461,8 +461,13 @@ def add_header(self, header): """Add headers to the email globaly or to a specific Personalization :param value: A Header object or a dict of header key/values + If dict has multiple keys, only first key/value item are applied :type value: Header, dict """ + if isinstance(header, dict): + (k, v) = list(header.items())[0] + self._headers = self._ensure_append(Header(k, v), self._headers) + return if header.personalization is not None: try: personalization = \ @@ -471,23 +476,14 @@ def add_header(self, header): except IndexError: personalization = Personalization() has_internal_personalization = False - if isinstance(header, dict): - (k, v) = list(header.items())[0] - personalization.add_header(Header(k, v)) - else: - personalization.add_header(header) + personalization.add_header(header) if not has_internal_personalization: self.add_personalization( personalization, index=header.personalization) else: - if isinstance(header, dict): - (k, v) = list(header.items())[0] - self._headers = self._ensure_append( - Header(k, v), self._headers) - else: - self._headers = self._ensure_append(header, self._headers) + self._headers = self._ensure_append(header, self._headers) @property def substitution(self): diff --git a/test/unit/test_mail_helpers.py b/test/unit/test_mail_helpers.py index a7d08d89..67359c31 100644 --- a/test/unit/test_mail_helpers.py +++ b/test/unit/test_mail_helpers.py @@ -2,6 +2,8 @@ import json import unittest +from sendgrid.helpers.mail.header import Header + try: from email.message import EmailMessage except ImportError: @@ -338,6 +340,142 @@ def test_single_email_with_all_three_email_contents_to_single_recipient(self): json.loads(response_content_with_all_three_mime_contents) ) + def test_single_email_with_custom_header_globaly(self): + from sendgrid.helpers.mail import (Mail, From, To, Subject, + Header, PlainTextContent) + self.maxDiff = None + message = Mail( + from_email=From('test+from@example.com', 'Example From Name'), + to_emails=To('test+to@example.com', 'Example To Name'), + subject=Subject('Sending with SendGrid is Fun'), + plain_text_content=PlainTextContent( + 'and easy to do anywhere, even with Python') + ) + + message.header = Header('X-Test1', 'Test1') + response_content = json.dumps({ + "content": [ + { + "type": "text/plain", + "value": "and easy to do anywhere, even with Python" + } + ], + "from": { + "email": "test+from@example.com", + "name": "Example From Name" + }, + "personalizations": [ + { + "to": [ + { + "email": "test+to@example.com", + "name": "Example To Name" + } + ] + } + ], + "subject": "Sending with SendGrid is Fun", + "headers": { + "X-Test1": "Test1" + } + }) + self.assertEqual( + message.get(), + json.loads(response_content) + ) + + def test_single_email_with_custom_header_globaly_single_key_value_dict(self): + from sendgrid.helpers.mail import (Mail, From, To, Subject, + Header, PlainTextContent) + self.maxDiff = None + message = Mail( + from_email=From('test+from@example.com', 'Example From Name'), + to_emails=To('test+to@example.com', 'Example To Name'), + subject=Subject('Sending with SendGrid is Fun'), + plain_text_content=PlainTextContent( + 'and easy to do anywhere, even with Python') + ) + message.add_header({'X-Test1': 'Test1'}) + + response_content = json.dumps({ + "content": [ + { + "type": "text/plain", + "value": "and easy to do anywhere, even with Python" + } + ], + "from": { + "email": "test+from@example.com", + "name": "Example From Name" + }, + "personalizations": [ + { + "to": [ + { + "email": "test+to@example.com", + "name": "Example To Name" + } + ] + } + ], + "subject": "Sending with SendGrid is Fun", + "headers": { + "X-Test1": "Test1" + } + }) + self.assertEqual( + message.get(), + json.loads(response_content) + ) + + def test_single_email_with_custom_header_globaly_multi_key_values_dict(self): + from sendgrid.helpers.mail import (Mail, From, To, Subject, + Header, PlainTextContent) + self.maxDiff = None + message = Mail( + from_email=From('test+from@example.com', 'Example From Name'), + to_emails=To('test+to@example.com', 'Example To Name'), + subject=Subject('Sending with SendGrid is Fun'), + plain_text_content=PlainTextContent( + 'and easy to do anywhere, even with Python') + ) + message.add_header({ + 'X-Test1': 'Test1', + 'X-Test2': 'Test2' + }) + + response_content = json.dumps({ + "content": [ + { + "type": "text/plain", + "value": "and easy to do anywhere, even with Python" + } + ], + "from": { + "email": "test+from@example.com", + "name": "Example From Name" + }, + "personalizations": [ + { + "to": [ + { + "email": "test+to@example.com", + "name": "Example To Name" + } + ] + } + ], + "subject": "Sending with SendGrid is Fun", + "headers": { + "X-Test1": "Test1", + # "X-Test2": "Test2", # does NOT applied + } + }) + self.assertEqual( + message.get(), + json.loads(response_content) + ) + def test_single_email_with_amp_and_html_contents_to_single_recipient(self): from sendgrid.helpers.mail import (Mail, From, To, Subject, PlainTextContent, HtmlContent, AmpHtmlContent)