Fix add_or_replace_header: make sure to delete all header occurrences before

This commit is contained in:
Son NK 2020-02-07 22:02:33 +07:00
parent 4c383663f4
commit dec36ceeec
2 changed files with 16 additions and 5 deletions

View File

@ -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):

View File

@ -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")]