瀏覽代碼

add: allow headers to be overridden

Mikael Koli 3 年之前
父節點
當前提交
d7e98cbdfe
共有 2 個文件被更改,包括 77 次插入13 次删除
  1. 20 13
      redmail/email/sender.py
  2. 57 0
      redmail/test/email/test_headers.py

+ 20 - 13
redmail/email/sender.py

@@ -425,26 +425,33 @@ class EmailSender:
 
     def _create_body(self, subject, sender, receivers=None, cc=None, bcc=None, headers=None) -> EmailMessage:
         msg = EmailMessage()
-        msg["From"] = sender
-        msg["Subject"] = subject
-        
+
+        email_headers = {
+            "From": sender,
+            "Subject": subject,
+        }
+
         # To whoom the email goes
         if receivers:
-            msg["To"] = receivers
+            email_headers["To"] = receivers
         if cc:
-            msg['Cc'] = cc
+            email_headers['Cc'] = cc
         if bcc:
-            msg['Bcc'] = bcc
+            email_headers['Bcc'] = bcc
+
+        email_headers.update({
+            # Message-IDs could be produced by the first mail server
+            # or the program sending the email (as we are doing now).
+            # Apparently Gmail might require it as of 2022
+            "Message-ID": self.create_message_id(),
 
-        # Message-IDs could be produced by the first mail server
-        # or the program sending the email (as we are doing now).
-        # Apparently Gmail might require it as of 2022
-        msg['Message-ID'] = self.create_message_id()
-        msg['Date'] = formatdate()
+            "Date": formatdate(),
+        })
 
         if headers:
-            for key, val in headers.items():
-                msg[key] = val
+            email_headers.update(headers)
+        for key, val in email_headers.items():
+            msg[key] = val
         return msg
 
     def _set_content_type(self, msg:EmailMessage, has_text, has_html, has_attachments):

+ 57 - 0
redmail/test/email/test_headers.py

@@ -4,6 +4,8 @@ from textwrap import dedent
 import sys
 import re
 
+import pytest
+
 from redmail import EmailSender
 
 from convert import remove_email_content_id, prune_generated_headers
@@ -72,3 +74,58 @@ def test_cc_bcc():
     Date: <date>
 
     """)[1:]
+
+@pytest.mark.parametrize("how", ["instance", "email"])
+def test_custom_headers(how):
+    email = EmailSender(host=None, port=1234)
+    headers = {"Importance": "high"}
+
+    if how == "email":
+        msg = email.get_message(
+            sender="me@example.com", 
+            subject="Some email", 
+            headers=headers
+        )
+    elif how == "instance":
+        email.headers = headers
+        msg = email.get_message(
+            sender="me@example.com", 
+            subject="Some email",
+        )
+    msg = prune_generated_headers(str(msg))
+    assert remove_email_content_id(msg) == dedent("""
+    From: me@example.com
+    Subject: Some email
+    Message-ID: <<message_id>>
+    Date: <date>
+    Importance: high
+
+    """)[1:]
+
+@pytest.mark.parametrize("how", ["instance", "email"])
+def test_custom_headers_override(how):
+    email = EmailSender(host=None, port=1234)
+    headers = {
+        "Date": "Sun, 31 Jan 2021 06:56:46 -0000",
+        "Message-ID": "<167294165062.31860.1664530310632362057@LAPTOP-1234GML0>"
+    }
+
+    if how == "email":
+        msg = email.get_message(
+            sender="me@example.com", 
+            subject="Some email",
+            headers=headers
+        )
+    elif how == "instance":
+        email.headers = headers
+        msg = email.get_message(
+            sender="me@example.com",
+            subject="Some email",
+        )
+    assert str(msg) == dedent("""
+    From: me@example.com
+    Subject: Some email
+    Message-ID: <167294165062.31860.1664530310632362057@LAPTOP-1234GML0>
+    Date: Sun, 31 Jan 2021 06:56:46 -0000
+
+    """)[1:]