homelab-codex-ws/jobs/gmail-bulk-import/tests/test_importer.py

"""Unit tests for Gmail bulk importer — no DB, no external services required."""
from __future__ import annotations

import hashlib
import json
import mailbox
from datetime import datetime, timezone
from email import encoders
from email.mime.base import MIMEBase
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from pathlib import Path

import pytest

from gmail_bulk_import import importer
from gmail_bulk_import.importer import (
    _message_id,
    _parse_attachments,
    _parse_date,
    run_import,
)


def _make_msg(headers: dict) -> mailbox.mboxMessage:
    msg = mailbox.mboxMessage()
    for k, v in headers.items():
        msg[k] = v
    return msg


def _make_mbox(path: Path, msgs: list[mailbox.mboxMessage]) -> None:
    mbox = mailbox.mbox(str(path), create=True)
    for msg in msgs:
        mbox.add(msg)
    mbox.flush()
    mbox.close()


def _msg_with_id(mid: str = "<test-unique@example.com>") -> mailbox.mboxMessage:
    return _make_msg({
        "From": "alice@example.com",
        "Message-ID": mid,
        "Date": "Tue, 10 Jun 2025 12:00:00 +0000",
        "Subject": "Hello",
    })


def _mime_with_attachments(
    mid: str = "<attach-test@example.com>",
    payloads: list | None = None,
) -> mailbox.mboxMessage:
    """Build a multipart mboxMessage with binary attachments.

    payloads: list of (filename, content_type, bytes). Defaults to PDF + PNG.
    """
    if payloads is None:
        payloads = [
            ("report.pdf", "application/pdf", b"PDF content here"),
            ("photo.png", "image/png", b"PNG bytes"),
        ]
    mime = MIMEMultipart()
    mime["From"] = "alice@example.com"
    mime["To"] = "bob@example.com"
    mime["Subject"] = "Files attached"
    mime["Message-ID"] = mid
    mime["Date"] = "Tue, 10 Jun 2025 12:00:00 +0000"
    mime.attach(MIMEText("See attachments", "plain"))
    for filename, ctype, data in payloads:
        maintype, subtype = ctype.split("/", 1)
        att = MIMEBase(maintype, subtype)
        att.set_payload(data)
        encoders.encode_base64(att)
        att.add_header("Content-Disposition", "attachment", filename=filename)
        mime.attach(att)
    return mailbox.mboxMessage(mime)


class TestMessageId:
    def test_uses_message_id_header(self):
        msg = _make_msg({"Message-ID": "<foo-123@example.com>"})
        assert _message_id(msg) == "foo-123@example.com"

    def test_strips_angle_brackets(self):
        msg = _make_msg({"Message-ID": "  <bar@baz.com>  "})
        assert _message_id(msg) == "bar@baz.com"

    def test_fallback_to_sha256_when_missing(self):
        msg = _make_msg({"From": "a@b.com", "Subject": "no id"})
        result = _message_id(msg)
        assert result.startswith("sha256-")
        assert len(result) == len("sha256-") + 32

    def test_sha256_is_stable(self):
        msg = _make_msg({"From": "a@b.com", "Subject": "stable"})
        assert _message_id(msg) == _message_id(msg)


class TestParseDate:
    def test_parses_rfc2822(self):
        msg = _make_msg({"Date": "Tue, 10 Jun 2025 12:00:00 +0000"})
        ts = _parse_date(msg)
        assert ts == datetime(2025, 6, 10, 12, 0, 0, tzinfo=timezone.utc)

    def test_result_is_utc_aware(self):
        msg = _make_msg({"Date": "Tue, 10 Jun 2025 12:00:00 +0000"})
        assert _parse_date(msg).tzinfo is not None

    def test_falls_back_to_epoch_on_bad_date(self):
        msg = _make_msg({"Date": "not-a-date"})
        assert _parse_date(msg) == datetime(1970, 1, 1, tzinfo=timezone.utc)

    def test_falls_back_to_epoch_when_missing(self):
        msg = _make_msg({"From": "a@b.com"})
        assert _parse_date(msg) == datetime(1970, 1, 1, tzinfo=timezone.utc)

    def test_converts_offset_to_utc(self):
        msg = _make_msg({"Date": "Tue, 10 Jun 2025 14:00:00 +0200"})
        ts = _parse_date(msg)
        assert ts == datetime(2025, 6, 10, 12, 0, 0, tzinfo=timezone.utc)


class TestParseAttachments:
    def test_plain_text_message_has_no_attachments(self):
        assert _parse_attachments(_msg_with_id()) == []

    def test_two_attachments_manifest(self):
        atts = _parse_attachments(_mime_with_attachments())
        assert len(atts) == 2

        pdf = next(a for a in atts if a["filename"] == "report.pdf")
        assert pdf["type"] == "attachment"
        assert pdf["content_type"] == "application/pdf"
        assert pdf["size"] == len(b"PDF content here")
        assert pdf["sha256"] == hashlib.sha256(b"PDF content here").hexdigest()

        png = next(a for a in atts if a["filename"] == "photo.png")
        assert png["content_type"] == "image/png"
        assert png["size"] == len(b"PNG bytes")
        assert png["sha256"] == hashlib.sha256(b"PNG bytes").hexdigest()

    def test_attachment_descriptors_have_required_keys(self):
        for att in _parse_attachments(_mime_with_attachments()):
            assert {"type", "filename", "content_type", "size", "sha256"} <= att.keys()


class TestRunImport:
    async def test_dry_run_counts_without_writing(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id("<a@e.com>"), _msg_with_id("<b@e.com>")])

        stats = await run_import(mbox_path, archive, dry_run=True)

        assert stats["processed"] == 2
        assert stats["imported"] == 2
        assert not archive.exists()

    async def test_imports_eml_to_archive(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id()])

        stats = await run_import(mbox_path, archive)

        assert stats["imported"] == 1
        assert stats["skipped"] == 0
        assert stats["errors"] == 0
        assert len(list(archive.rglob("*.eml"))) == 1

    async def test_eml_stored_under_gmail_source(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id()])

        await run_import(mbox_path, archive)

        eml_files = list(archive.rglob("*.eml"))
        assert eml_files[0].parts[len(archive.parts)] == "gmail"

    async def test_idempotent_second_run_skips(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id()])

        await run_import(mbox_path, archive)
        stats2 = await run_import(mbox_path, archive)

        assert stats2["imported"] == 0
        assert stats2["skipped"] == 1
        assert stats2["errors"] == 0

    async def test_two_distinct_messages_both_imported(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id("<msg1@e.com>"), _msg_with_id("<msg2@e.com>")])

        stats = await run_import(mbox_path, archive)

        assert stats["imported"] == 2
        assert stats["errors"] == 0
        assert len(list(archive.rglob("*.eml"))) == 2

    async def test_archive_only_mode_no_dsn(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id()])

        stats = await run_import(mbox_path, archive, dsn=None)

        assert stats["imported"] == 1
        assert list(archive.rglob("*.eml"))

    async def test_limit_truncates_processing(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(5)])

        stats = await run_import(mbox_path, archive, limit=3)

        assert stats["processed"] == 3
        assert stats["imported"] == 3

    async def test_limit_with_dry_run(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(10)])

        stats = await run_import(mbox_path, archive, dry_run=True, limit=4)

        assert stats["processed"] == 4
        assert stats["imported"] == 4
        assert not archive.exists()

    async def test_epoch_fallback_counted(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [
            _make_msg({"Message-ID": "<bad@e.com>", "Date": "not-a-date"}),
            _make_msg({"Message-ID": "<missing@e.com>"}),
            _msg_with_id("<good@e.com>"),
        ])

        stats = await run_import(mbox_path, archive)

        assert stats["epoch_fallback"] == 2
        assert stats["processed"] == 3

    async def test_attachment_stats_aggregated(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [
            _msg_with_id("<plain@e.com>"),
            _mime_with_attachments("<att@e.com>"),
        ])

        stats = await run_import(mbox_path, archive)

        assert stats["msgs_with_attachments"] == 1
        assert stats["total_attachments"] == 2
        assert stats["total_attachment_bytes"] == (
            len(b"PDF content here") + len(b"PNG bytes")
        )

    async def test_batch_flush_multiple_messages(self, tmp_path):
        """BATCH_SIZE=500; 3 messages stay in pending until end-of-loop flush."""
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(3)])

        stats = await run_import(mbox_path, archive)

        assert stats["imported"] == 3
        assert stats["errors"] == 0
        assert len(list(archive.rglob("*.eml"))) == 3

    async def test_attachment_entities_in_stats(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_mime_with_attachments()])

        stats = await run_import(mbox_path, archive)

        assert stats["imported"] == 1
        assert stats["total_attachments"] == 2


# ---------------------------------------------------------------------------
# Regression tests for the 2026-07-14 hardening audit (four fixes).
#
# The audit reproduced each of these on crafted mboxes; the corpus is proven to
# contain raw 8-bit header bytes, so every scenario below is a real re-import
# risk, not a hypothetical.
# ---------------------------------------------------------------------------

# Raw 8-bit byte (0xe9) inside Message-ID: compat32 .get() returns an
# email.header.Header, whose .strip() raised AttributeError. Because
# _message_id ran BEFORE the per-message try, one such header killed the run.
_EIGHTBIT_MID_EML = (
    b"From: alice@example.com\r\n"
    b"Message-ID: <bad\xe9id@example.com>\r\n"
    b"Date: Tue, 10 Jun 2025 12:00:00 +0000\r\n"
    b"Subject: hi\r\n\r\nbody"
)

# Raw 8-bit byte inside the Date header's trailing comment: .get() returns a
# Header, parsedate_to_datetime(Header) raises → needless epoch_fallback, even
# though str(header) parses to a real date (2009-05-19 10:27:09 +0200).
_EIGHTBIT_DATE_EML = (
    b"From: alice@example.com\r\n"
    b"Message-ID: <clean-date@example.com>\r\n"
    b"Date: Tue, 19 May 2009 10:27:09 +0200 (Ho\xe9ra)\r\n"
    b"Subject: hi\r\n\r\nbody"
)


def _make_mbox_raw(path: Path, raw_messages: list[bytes]) -> None:
    mbox = mailbox.mbox(str(path), create=True)
    for raw in raw_messages:
        mbox.add(mailbox.mboxMessage(raw))
    mbox.flush()
    mbox.close()


class _FakeConn:
    """Stand-in asyncpg connection: records executemany rows or fails them."""

    def __init__(self, fail: bool = False):
        self.fail = fail
        self.inserted: list = []
        self.executemany_calls = 0

    async def executemany(self, query, rows):
        self.executemany_calls += 1
        if self.fail:
            raise RuntimeError("DB down")
        self.inserted.extend(list(rows))

    async def close(self):
        pass


def _patch_connect(monkeypatch, conn: _FakeConn) -> None:
    async def _fake_connect(dsn):
        return conn
    monkeypatch.setattr(importer.asyncpg, "connect", _fake_connect)


class TestEightBitMessageId:
    def test_message_id_8bit_header_does_not_crash(self):
        # Fix 1: str() + _sanitize gives .strip() a real string. Old code did
        # Header.strip() → AttributeError.
        msg = mailbox.mboxMessage(_EIGHTBIT_MID_EML)
        mid = _message_id(msg)
        assert isinstance(mid, str)
        assert mid.endswith("id@example.com")
        json.dumps(mid).encode("utf-8")  # jsonb-safe: no lone surrogate

    async def test_run_survives_8bit_message_id(self, tmp_path):
        # Old code raised out of run_import here; the run must now complete and
        # count the message normally.
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox_raw(mbox_path, [_EIGHTBIT_MID_EML, _msg_with_id("<ok@e.com>").as_bytes()])

        stats = await run_import(mbox_path, archive)

        assert stats["processed"] == 2
        assert stats["imported"] == 2
        assert stats["errors"] == 0


class TestPerMessageGuard:
    async def test_one_bad_message_counts_once_and_run_continues(self, tmp_path, monkeypatch):
        # Fix 2: the parse helpers now live inside the per-message try, so a
        # message that raises is one error, not the death of the run.
        calls = {"n": 0}
        real_message_id = importer._message_id

        def flaky(msg):
            calls["n"] += 1
            if calls["n"] == 1:
                raise RuntimeError("boom parsing first message")
            return real_message_id(msg)

        monkeypatch.setattr(importer, "_message_id", flaky)

        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id("<a@e.com>"), _msg_with_id("<b@e.com>")])

        stats = await run_import(mbox_path, archive)

        assert stats["processed"] == 2
        assert stats["errors"] == 1
        assert stats["imported"] == 1


class TestFailedInsertDoesNotPoison:
    async def test_final_flush_failure_does_not_kill_run(self, tmp_path, monkeypatch):
        # Fix 3: the final flush used to be uncaught (try/finally, no except),
        # so a failing insert lost all stats. Now it is swallowed and counted.
        conn = _FakeConn(fail=True)
        _patch_connect(monkeypatch, conn)

        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(3)])

        stats = await run_import(mbox_path, archive, dsn="postgresql://fake")

        # Run completes, archive is ahead of DB, and the drift is visible.
        assert stats["imported"] == 3
        assert stats["db_inserted"] == 0
        assert stats["db_insert_failed"] == 3
        assert stats["errors"] == 0

    async def test_poisoned_batch_does_not_re_flush_forever(self, tmp_path, monkeypatch):
        # Fix 3: pending is cleared even on insert failure, so a poisoned batch
        # is not re-flushed by every subsequent message. With BATCH_SIZE=1 each
        # message flushes exactly once; old code re-flushed a growing pending
        # list and died in the uncaught final flush.
        monkeypatch.setattr(importer, "BATCH_SIZE", 1)
        conn = _FakeConn(fail=True)
        _patch_connect(monkeypatch, conn)

        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(3)])

        stats = await run_import(mbox_path, archive, dsn="postgresql://fake")

        assert conn.executemany_calls == 3  # one per message, not a growing batch
        assert stats["imported"] == 3
        assert stats["db_insert_failed"] == 3
        assert stats["errors"] == 0


class TestStatsBalance:
    async def test_processed_balances_and_db_reconciles(self, tmp_path, monkeypatch):
        # Fix 4: imported/skipped count archive writes; db_inserted counts DB
        # rows. Both invariants must always hold.
        conn = _FakeConn()
        _patch_connect(monkeypatch, conn)
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(3)])

        stats1 = await run_import(mbox_path, archive, dsn="postgresql://fake")

        # second run over the same archive: every message is skipped, but each
        # still goes to the DB (ON CONFLICT DO NOTHING)
        conn2 = _FakeConn()
        _patch_connect(monkeypatch, conn2)
        stats2 = await run_import(mbox_path, archive, dsn="postgresql://fake")

        for s in (stats1, stats2):
            assert s["processed"] == s["imported"] + s["skipped"] + s["errors"]
            assert s["db_inserted"] + s["db_insert_failed"] == s["imported"] + s["skipped"]

        assert stats1["imported"] == 3 and stats1["db_inserted"] == 3
        assert stats2["skipped"] == 3 and stats2["db_inserted"] == 3

    async def test_archive_only_mode_leaves_db_counters_zero(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(2)])

        stats = await run_import(mbox_path, archive, dsn=None)

        assert stats["imported"] == 2
        assert stats["db_inserted"] == 0
        assert stats["db_insert_failed"] == 0
        assert stats["processed"] == stats["imported"] + stats["skipped"] + stats["errors"]


class TestParseDate8Bit:
    def test_8bit_date_header_parses_via_str_not_epoch(self):
        # Fix 4 (date): str(header) before epoch_fallback. Old code fed the
        # Header straight to parsedate_to_datetime → AttributeError → epoch.
        msg = mailbox.mboxMessage(_EIGHTBIT_DATE_EML)
        ts = _parse_date(msg)
        assert ts == datetime(2009, 5, 19, 8, 27, 9, tzinfo=timezone.utc)

    async def test_8bit_date_not_counted_as_epoch_fallback(self, tmp_path):
        mbox_path = tmp_path / "mail.mbox"
        archive = tmp_path / "archive"
        _make_mbox_raw(mbox_path, [_EIGHTBIT_DATE_EML])

        stats = await run_import(mbox_path, archive)

        assert stats["epoch_fallback"] == 0
        assert stats["imported"] == 1