rebuild test_smtp

Author: bmos

test/test_smtp/test_smtp.py

@@ -1,138 +1,196 @@
-import base64
-import io
-import re
-import unittest
+import random
+import string
+from email import message_from_string
+from email.message import Message
+from io import BytesIO, StringIO
+from unittest.mock import MagicMock
+
+import pytest
+from pytest_mock import MockerFixture
 
 from parsons import SMTP
+from parsons.notifications.sendmail import EmptyListError
+
+
+@pytest.fixture
+def mock_conn(mocker: MockerFixture) -> MagicMock:
+    mock_smtp_class = mocker.patch("parsons.notifications.smtp.smtplib.SMTP", autospec=True)
+    conn = mock_smtp_class.return_value
+    conn.sendmail.return_value = None
+
+    def get_msg() -> Message | None:
+        if not conn.sendmail.called:
+            return None
+        args, _ = conn.sendmail.call_args
+        return message_from_string(args[2])
+
+    def get_types() -> list[str]:
+        msg = get_msg()
+        if not msg:
+            return []
+        return [p.get_content_type() for p in msg.walk() if not p.is_multipart()]
+
+    conn.get_sent_msg = get_msg
+    conn.get_sent_types = get_types
+    return conn
+
+
+@pytest.fixture
+def smtp(mock_conn: MagicMock) -> SMTP:
+    """Configure a Parsons SMTP instance with fake credentials."""
+    return SMTP("fake.example.com", username="fake", password="fake")
+
+
+def create_virtual_file(name: str, content: str = "data") -> StringIO:
+    f = StringIO(content)
+    f.name = name
+    return f
+
+
+@pytest.mark.parametrize(
+    ("html", "files", "expected_types"),
+    [
+        (None, None, ["text/plain"]),
+        ("<p>HTML</p>", None, ["text/plain", "text/html"]),
+        (None, ["test.csv"], ["text/plain", "text/csv"]),
+        ("<p>HTML</p>", ["test.csv"], ["text/plain", "text/html", "text/csv"]),
+        (
+            None,
+            ["1.csv", "2.csv"],
+            ["text/plain", "text/csv", "text/csv"],
+        ),
+    ],
+    ids=(
+        "Simple Body",
+        "Body + HTML",
+        "Body + Attachment",
+        "Body + HTML + Attachment",
+        "Multiple Attachments",
+    ),
+)
+def test_send_email_content(
+    smtp: SMTP,
+    mock_conn: MagicMock,
+    html: str | None,
+    files: list[str] | None,
+    expected_types: list[str],
+    mocker: MockerFixture,
+):
+    files_with_data = (
+        [
+            create_virtual_file(
+                f, content="".join(random.choices(string.ascii_letters + string.digits, k=8))
+            )
+            for f in files
+        ]
+        if files
+        else None
+    )
+
+    smtp.send_email(
+        "f@ex.com",
+        "t@ex.com",
+        "We've been trying to reach you...",
+        "Your car's warranty is about to expire!",
+        message_html=html,
+        files=files_with_data,
+    )
+
+    msg = mock_conn.get_sent_msg()
+    assert msg["Subject"] == "We've been trying to reach you..."
+    assert msg["To"] == "t@ex.com"
+    assert msg["From"] == "f@ex.com"
+    sent_payloads = [p.get_payload(decode=True) for p in msg.walk() if not p.is_multipart()]
+    assert b"Your car's warranty is about to expire!" in sent_payloads
+
+    if html:
+        assert html.encode() in sent_payloads
+
+    if files_with_data:
+        for content in files_with_data:
+            assert content.read().encode() in sent_payloads
+
+    sent_types = mock_conn.get_sent_types()
+    for t in expected_types:
+        assert t in sent_types
+
+
+def test_binary_attachment_integrity(smtp: SMTP, mock_conn: MagicMock):
+    gif_data = bytes([71, 73, 70, 56, 57, 97, 1, 0, 1, 0])
+    gif_file = BytesIO(gif_data)
+    gif_file.name = "tiny.gif"
+
+    smtp.send_email("f@ex.com", "t@ex.com", "Sub", "Body", files=[gif_file])
+
+    msg = mock_conn.get_sent_msg()
+    gif_part = next(p for p in msg.walk() if p.get_filename() == "tiny.gif")
+
+    assert gif_part.get_content_type() == "image/gif"
+    assert gif_part.get_payload(decode=True) == gif_data
+
+
+def test_attachment_disposition(smtp: SMTP, mock_conn: MagicMock):
+    named_file = StringIO("content")
+    named_file.name = "report.pdf"
+
+    smtp.send_email("f@ex.com", "t@ex.com", "Sub", "Body", files=[named_file])
+
+    msg = mock_conn.get_sent_msg()
+    file_part = next(p for p in msg.walk() if p.get_filename() == "report.pdf")
+    assert "attachment" in file_part.get("Content-Disposition")
+
+
+def test_send_email_files_as_single_string(smtp: SMTP, mock_conn: MagicMock, mocker: MockerFixture):
+    # We mock the attachment creator because we only care about the input conversion here
+    mock_create = mocker.patch.object(smtp, "_create_message_attachments")
+    filename = "single_report.pdf"
+
+    smtp.send_email("sender@ex.com", "to@ex.com", "Subject", "Body", files=filename)
+
+    args, _ = mock_create.call_args
+    assert args[4] == [filename]
+
+
+def test_send_email_empty_recipient_list(smtp: SMTP):
+    with pytest.raises(EmptyListError, match="Must contain at least 1 email."):
+        smtp.send_email("sender@ex.com", [], "Sub", "Body")
+
+
+@pytest.mark.parametrize(
+    ("input_to", "expected_header"),
+    [
+        (["a@b.com", "c@d.com"], "a@b.com, c@d.com"),
+        (['"Name" <a@b.com>', "c@d.com"], '"Name" <a@b.com>, c@d.com'),
+    ],
+)
+def test_recipient_string_joining(
+    smtp: SMTP, mock_conn: MagicMock, input_to: list[str], expected_header: str
+):
+    smtp.send_email("f@ex.com", input_to, "Sub", "Body")
+
+    msg = mock_conn.get_sent_msg()
+    assert msg["To"] == expected_header
+
+
+@pytest.mark.parametrize(("close_manually", "expected_quit_count"), [(True, 0), (False, 1)])
+def test_connection_closing_logic(
+    mock_conn: MagicMock, close_manually: bool, expected_quit_count: int
+) -> None:
+    smtp_inst = SMTP(
+        "fake.example.com", username="fake", password="fake", close_manually=close_manually
+    )
+
+    smtp_inst.send_email("a@b.com", "c@d.com", "Sub", "Body")
+
+    assert mock_conn.quit.call_count == expected_quit_count
+
+
+def test_send_message_error_handling(smtp: SMTP, mock_conn: MagicMock) -> None:
+    """Ensure partial failures are returned correctly."""
+    fail_response = {"bad@ex.com": (550, "User unknown")}
+    mock_conn.sendmail.return_value = fail_response
 
+    simple_msg = smtp._create_message_simple("f@ex.com", "bad@ex.com", "Sub", "Body")
+    result = smtp._send_message(simple_msg)
 
-class FakeConnection:
-    def __init__(self, result_obj):
-        self.result_obj = result_obj
-
-    def sendmail(self, sender, to, message_body):
-        self.result_obj.result = (sender, to, message_body)
-        if "willfail@example.com" in to:
-            return {"willfail@example.com": (550, "User unknown")}
-
-    def quit(self):
-        self.result_obj.quit_ran = True
-
-
-class TestSMTP(unittest.TestCase):
-    def setUp(self):
-        self.smtp = SMTP("fake.example.com", username="fake", password="fake")
-        self.smtp.conn = FakeConnection(self)
-        self.result = None
-        self.quit_ran = False
-
-    def test_send_message_simple(self):
-        self.smtp.send_email(
-            "foo@example.com", "recipient1@example.com", "Simple subject", "Fake body"
-        )
-        assert self.result[0] == "foo@example.com"
-        assert self.result[1] == ["recipient1@example.com"]
-        assert self.result[2].endswith(
-            "\nto: recipient1@example.com\nfrom: foo@example.com"
-            "\nsubject: Simple subject\n\nFake body"
-        )
-        assert self.quit_ran
-
-    def test_send_message_html(self):
-        self.smtp.send_email(
-            "foohtml@example.com",
-            "recipienthtml@example.com",
-            "Simple subject",
-            "Fake body",
-            "<p>Really Fake html</p>",
-        )
-        assert self.result[0] == "foohtml@example.com"
-        assert self.result[1] == ["recipienthtml@example.com"]
-        assert re.search(r"<p>Really Fake html</p>\n--=======", self.result[2])
-        assert re.search(r"\nFake body\n--======", self.result[2])
-        assert re.search(r"ubject: Simple subject\n", self.result[2])
-        assert self.quit_ran
-
-    def test_send_message_manualclose(self):
-        smtp = SMTP("fake.example.com", username="fake", password="fake", close_manually=True)
-        smtp.conn = FakeConnection(self)
-        smtp.send_email("foo@example.com", "recipient1@example.com", "Simple subject", "Fake body")
-        assert not self.quit_ran
-
-    def test_send_message_files(self):
-        named_file_content = "x,y,z\n1,2,3\r\n3,4,5\r\n"
-        unnamed_file_content = "foo,bar\n1,2\r\n3,4\r\n"
-        bytes_file_content = bytes(
-            [
-                71,
-                73,
-                70,
-                56,
-                57,
-                97,
-                1,
-                0,
-                1,
-                0,
-                0,
-                255,
-                0,
-                44,
-                0,
-                0,
-                0,
-                0,
-                1,
-                0,
-                1,
-                0,
-                0,
-                2,
-                0,
-                59,
-            ]
-        )
-        named_file = io.StringIO(named_file_content)
-        named_file.name = "xyz.csv"
-
-        bytes_file = io.BytesIO(bytes_file_content)
-        bytes_file.name = "xyz.gif"
-
-        self.smtp.send_email(
-            "foofiles@example.com",
-            "recipientfiles@example.com",
-            "Simple subject",
-            "Fake body",
-            files=[io.StringIO(unnamed_file_content), named_file, bytes_file],
-        )
-        assert self.result[0] == "foofiles@example.com"
-        assert self.result[1] == ["recipientfiles@example.com"]
-        assert re.search(r"\nFake body\n--======", self.result[2])
-        found = re.findall(r'filename="file"\n\n([\w=/]+)\n\n--===', self.result[2])
-        assert base64.b64decode(found[0]).decode() == unnamed_file_content
-        found_named = re.findall(
-            r'Content-Type: text/csv; charset="utf-8"\nMIME-Version: 1.0'
-            r"\nContent-Transfer-Encoding: base64\nContent-Disposition: "
-            r'attachment; filename="xyz.csv"\n\n([\w=/]+)\n\n--======',
-            self.result[2],
-        )
-        assert base64.b64decode(found_named[0]).decode() == named_file_content
-
-        found_gif = re.findall(
-            r"Content-Type: image/gif\nMIME-Version: 1.0"
-            r"\nContent-Transfer-Encoding: base64\nContent-ID: <xyz.gif>"
-            r'\nContent-Disposition: attachment; filename="xyz.gif"\n\n([\w=/]+)\n\n--==',
-            self.result[2],
-        )
-        assert base64.b64decode(found_gif[0]) == bytes_file_content
-        assert self.quit_ran
-
-    def test_send_message_partial_fail(self):
-        simple_msg = self.smtp._create_message_simple(
-            "foo@example.com",
-            "recipient1@example.com, willfail@example.com",
-            "Simple subject",
-            "Fake body",
-        )
-        send_result = self.smtp._send_message(simple_msg)
-        assert send_result == {"willfail@example.com": (550, "User unknown")}
+    assert result == fail_response