From dec36ceeecabea88fb09dc92c513373746e11c74 Mon Sep 17 00:00:00 2001 From: Son NK Date: Fri, 7 Feb 2020 22:02:33 +0700 Subject: [PATCH] Fix add_or_replace_header: make sure to delete all header occurrences before --- app/email_utils.py | 10 +++++----- tests/test_email_utils.py | 11 +++++++++++ 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/app/email_utils.py b/app/email_utils.py index f2833640..c635ccb3 100644 --- a/app/email_utils.py +++ b/app/email_utils.py @@ -290,11 +290,11 @@ def add_dkim_signature(msg: Message, email_domain: str): def add_or_replace_header(msg: Message, header: str, value: str): - try: - msg.add_header(header, value) - except ValueError: - # the header exists already - msg.replace_header(header, value) + """ + Remove all occurrences of `header` and add `header` with `value`. + """ + delete_header(msg, header) + msg[header] = value def delete_header(msg: Message, header: str): diff --git a/tests/test_email_utils.py b/tests/test_email_utils.py index 15d3b539..27e8b2a1 100644 --- a/tests/test_email_utils.py +++ b/tests/test_email_utils.py @@ -8,6 +8,7 @@ from app.email_utils import ( email_belongs_to_alias_domains, can_be_used_as_personal_email, delete_header, + add_or_replace_header, ) from app.extensions import db from app.models import User, CustomDomain @@ -71,3 +72,13 @@ def test_delete_header(): delete_header(msg, "H") assert msg._headers == [] + + +def test_add_or_replace_header(): + msg = EmailMessage() + msg["H"] = "abcd" + msg["H"] = "xyzt" + assert msg._headers == [("H", "abcd"), ("H", "xyzt")] + + add_or_replace_header(msg, "H", "new") + assert msg._headers == [("H", "new")]