stack/tests/bib/test_email_ingest_exercise.py

"""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 _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 = MagicMock()
        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 = MagicMock()
        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 = MagicMock()
        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 = MagicMock()
        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 = MagicMock()
        stats = ingest(store, _make_mailbox())
        assert stats["errors"] == 1


class TestIngestMessage:
    def test_basic(self, tmp_path):
        store = MagicMock()
        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 = MagicMock()
        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 = MagicMock()
        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 = MagicMock()
        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 = MagicMock()
        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("") == ""