229 lines
7.6 KiB
Python
229 lines
7.6 KiB
Python
"""Exercise bib.email_ingest — ingest() with mocked IMAP, _ingest_message with real msg."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import email
|
|
from email import encoders
|
|
from email.mime.base import MIMEBase
|
|
from email.mime.multipart import MIMEMultipart
|
|
from email.mime.text import MIMEText
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from bib.email_ingest import (
|
|
Mailbox,
|
|
_extract_body,
|
|
_ingest_message,
|
|
_sender_domain,
|
|
ingest,
|
|
)
|
|
|
|
|
|
def _mock_store() -> MagicMock:
|
|
"""MagicMock store safe to pass through _ingest_message.
|
|
|
|
_ingest_message's resend-dedup lookup calls
|
|
``store._con().execute(...).fetchone()`` (refs #665) — on a bare
|
|
MagicMock that returns a truthy Mock, which _ingest_message
|
|
misreads as an existing duplicate and short-circuits before ever
|
|
calling store.upsert(). Configuring fetchone() to return None
|
|
makes the mock behave like an empty store.
|
|
"""
|
|
s = MagicMock()
|
|
s._con.return_value.execute.return_value.fetchone.return_value = None
|
|
return s
|
|
|
|
|
|
def _make_mailbox():
|
|
return Mailbox(
|
|
host="mail.example.com",
|
|
port=993,
|
|
username="test@example.com",
|
|
password="pw",
|
|
folder="INBOX",
|
|
)
|
|
|
|
|
|
def _make_msg(
|
|
subject="Test Subject",
|
|
sender="user@cms.hhs.gov",
|
|
body="Hello world",
|
|
html=None,
|
|
attachment_name=None,
|
|
attachment_data=None,
|
|
list_id=None,
|
|
date="Thu, 17 Apr 2025 12:00:00 +0000",
|
|
message_id="<abc123@example.com>",
|
|
):
|
|
msg = MIMEMultipart()
|
|
msg["Subject"] = subject
|
|
msg["From"] = f"Sender <{sender}>"
|
|
msg["Date"] = date
|
|
msg["Message-ID"] = message_id
|
|
if list_id:
|
|
msg["List-ID"] = list_id
|
|
if body:
|
|
msg.attach(MIMEText(body, "plain"))
|
|
if html:
|
|
msg.attach(MIMEText(html, "html"))
|
|
if attachment_name and attachment_data:
|
|
part = MIMEBase("application", "octet-stream")
|
|
part.set_payload(attachment_data)
|
|
encoders.encode_base64(part)
|
|
part.add_header(
|
|
"Content-Disposition", f'attachment; filename="{attachment_name}"'
|
|
)
|
|
msg.attach(part)
|
|
return msg
|
|
|
|
|
|
class TestIngest:
|
|
@patch("bib.email_ingest.imaplib.IMAP4_SSL")
|
|
def test_basic_flow(self, mc_imap_cls):
|
|
conn = MagicMock()
|
|
mc_imap_cls.return_value = conn
|
|
conn.uid.side_effect = [
|
|
("OK", [b"1 2"]), # SEARCH
|
|
("OK", [(b"1", _make_msg().as_bytes())]), # FETCH uid 1
|
|
("OK", None), # STORE uid 1
|
|
("OK", [(b"2", _make_msg(subject="Second").as_bytes())]), # FETCH uid 2
|
|
("OK", None), # STORE uid 2
|
|
]
|
|
store = _mock_store()
|
|
store.upsert.return_value = "KEY1"
|
|
stats = ingest(store, _make_mailbox())
|
|
assert stats["seen"] == 2
|
|
assert stats["ingested"] == 2
|
|
|
|
@patch("bib.email_ingest.imaplib.IMAP4_SSL")
|
|
def test_search_fails(self, mc_imap_cls):
|
|
conn = MagicMock()
|
|
mc_imap_cls.return_value = conn
|
|
conn.uid.return_value = ("NO", [b""])
|
|
store = _mock_store()
|
|
stats = ingest(store, _make_mailbox())
|
|
assert stats["ingested"] == 0
|
|
|
|
@patch("bib.email_ingest.imaplib.IMAP4_SSL")
|
|
def test_fetch_fails(self, mc_imap_cls):
|
|
conn = MagicMock()
|
|
mc_imap_cls.return_value = conn
|
|
conn.uid.side_effect = [
|
|
("OK", [b"1"]), # SEARCH
|
|
("OK", [None]), # FETCH returns None
|
|
]
|
|
store = _mock_store()
|
|
stats = ingest(store, _make_mailbox())
|
|
assert stats["errors"] == 1
|
|
|
|
@patch("bib.email_ingest.imaplib.IMAP4_SSL")
|
|
def test_with_limit(self, mc_imap_cls):
|
|
conn = MagicMock()
|
|
mc_imap_cls.return_value = conn
|
|
conn.uid.side_effect = [
|
|
("OK", [b"1 2 3"]), # SEARCH
|
|
("OK", [(b"1", _make_msg().as_bytes())]), # FETCH uid 1
|
|
("OK", None), # STORE uid 1
|
|
]
|
|
store = _mock_store()
|
|
store.upsert.return_value = "KEY1"
|
|
stats = ingest(store, _make_mailbox(), limit=1)
|
|
assert stats["seen"] == 1
|
|
|
|
@patch("bib.email_ingest.imaplib.IMAP4_SSL")
|
|
def test_ingest_exception(self, mc_imap_cls):
|
|
conn = MagicMock()
|
|
mc_imap_cls.return_value = conn
|
|
conn.uid.side_effect = [
|
|
("OK", [b"1"]),
|
|
Exception("boom"),
|
|
]
|
|
store = _mock_store()
|
|
stats = ingest(store, _make_mailbox())
|
|
assert stats["errors"] == 1
|
|
|
|
|
|
class TestIngestMessage:
|
|
def test_basic(self, tmp_path):
|
|
store = _mock_store()
|
|
store.upsert.return_value = "KEY1"
|
|
msg = _make_msg()
|
|
parsed = email.message_from_bytes(msg.as_bytes(), policy=email.policy.default)
|
|
n = _ingest_message(store, _make_mailbox(), parsed, tmp_path)
|
|
assert n == 0
|
|
store.upsert.assert_called_once()
|
|
|
|
def test_with_attachment(self, tmp_path):
|
|
store = _mock_store()
|
|
store.upsert.return_value = "KEY1"
|
|
msg = _make_msg(attachment_name="doc.pdf", attachment_data=b"PDF content")
|
|
parsed = email.message_from_bytes(msg.as_bytes(), policy=email.policy.default)
|
|
n = _ingest_message(store, _make_mailbox(), parsed, tmp_path)
|
|
assert n == 1
|
|
|
|
def test_no_message_id(self, tmp_path):
|
|
store = _mock_store()
|
|
store.upsert.return_value = "KEY1"
|
|
msg = _make_msg(message_id="")
|
|
del msg["Message-ID"]
|
|
parsed = email.message_from_bytes(msg.as_bytes(), policy=email.policy.default)
|
|
n = _ingest_message(store, _make_mailbox(), parsed, tmp_path)
|
|
assert n == 0
|
|
|
|
def test_with_list_id(self, tmp_path):
|
|
store = _mock_store()
|
|
store.upsert.return_value = "KEY1"
|
|
msg = _make_msg(list_id="<cms-updates.listserv.cms.gov>")
|
|
parsed = email.message_from_bytes(msg.as_bytes(), policy=email.policy.default)
|
|
n = _ingest_message(store, _make_mailbox(), parsed, tmp_path)
|
|
assert n == 0
|
|
|
|
def test_unnamed_attachment(self, tmp_path):
|
|
store = _mock_store()
|
|
store.upsert.return_value = "KEY1"
|
|
msg = MIMEMultipart()
|
|
msg["Subject"] = "Test"
|
|
msg["From"] = "user@test.com"
|
|
msg["Date"] = "Thu, 17 Apr 2025 12:00:00 +0000"
|
|
msg["Message-ID"] = "<test@test.com>"
|
|
msg.attach(MIMEText("body", "plain"))
|
|
part = MIMEBase("application", "octet-stream")
|
|
part.set_payload(b"data")
|
|
encoders.encode_base64(part)
|
|
# No filename header
|
|
msg.attach(part)
|
|
parsed = email.message_from_bytes(msg.as_bytes(), policy=email.policy.default)
|
|
n = _ingest_message(store, _make_mailbox(), parsed, tmp_path)
|
|
assert n >= 0 # May or may not be treated as attachment
|
|
|
|
|
|
class TestExtractBody:
|
|
def test_plain_text(self):
|
|
msg = _make_msg(body="Hello plain")
|
|
parsed = email.message_from_bytes(msg.as_bytes(), policy=email.policy.default)
|
|
body = _extract_body(parsed)
|
|
assert "Hello plain" in body
|
|
|
|
def test_html_only(self):
|
|
msg = MIMEMultipart()
|
|
msg.attach(MIMEText("<p>Hello HTML</p>", "html"))
|
|
parsed = email.message_from_bytes(msg.as_bytes(), policy=email.policy.default)
|
|
body = _extract_body(parsed)
|
|
assert "Hello HTML" in body
|
|
|
|
def test_empty(self):
|
|
msg = MIMEMultipart()
|
|
parsed = email.message_from_bytes(msg.as_bytes(), policy=email.policy.default)
|
|
body = _extract_body(parsed)
|
|
assert body == ""
|
|
|
|
|
|
class TestSenderDomain:
|
|
def test_angle_bracket(self):
|
|
assert _sender_domain("User <user@cms.hhs.gov>") == "cms.hhs.gov"
|
|
|
|
def test_bare(self):
|
|
assert _sender_domain("user@cms.hhs.gov") == "cms.hhs.gov"
|
|
|
|
def test_empty(self):
|
|
assert _sender_domain("") == ""
|