webmail-client/backend/tests/indexer/test_pipeline.py

"""End-to-end indexer pipeline + Postgres search (SPEC §6.4, §8).

Uses a committed session (search reads through its own connection) and deletes
the user on teardown; the user FK cascade removes all created rows.
"""

from __future__ import annotations

import uuid
from collections.abc import AsyncIterator
from datetime import UTC, datetime
from email.message import EmailMessage

import pytest_asyncio
from sqlalchemy import delete, select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool

from app.config import get_settings
from app.indexer.indexer import persist_message
from app.models import Account, BodyPart, KnownRecipient, Mailbox, Message, User
from app.providers.base import RemoteMessage
from app.search.pg import PgSearchBackend
from app.search.queryparser import to_search_query


class Fixture:
    def __init__(
        self,
        session: AsyncSession,
        user_id: uuid.UUID,
        account_id: uuid.UUID,
        mailbox_id: uuid.UUID,
    ) -> None:
        self.session = session
        self.user_id = user_id
        self.account_id = account_id
        self.mailbox_id = mailbox_id
        self.backend = PgSearchBackend()
        self._uid = 0

    async def ingest(self, raw: bytes, *, flags: frozenset[str] = frozenset()) -> Message | None:
        self._uid += 1
        remote = RemoteMessage(uid=self._uid, raw=raw, flags=flags, internal_date=datetime.now(UTC))
        result = await persist_message(
            self.session,
            user_id=self.user_id,
            account_id=self.account_id,
            mailbox_id=self.mailbox_id,
            mailbox_special_use="inbox",
            remote=remote,
        )
        await self.session.commit()
        if result is not None:
            await self.backend.index(result.search_doc)
            return result.message
        return None


@pytest_asyncio.fixture
async def fx() -> AsyncIterator[Fixture]:
    engine = create_async_engine(get_settings().database_url, poolclass=NullPool)
    sm = async_sessionmaker(engine, expire_on_commit=False)
    session = sm()
    user = User(handle=f"pipe{uuid.uuid4().hex[:8]}@test.local", password_hash="x")
    session.add(user)
    await session.flush()
    account = Account(
        user_id=user.id,
        email="me@example.com",
        imap_config={},
        smtp_config={},
        secret_ct=b"",
        discovery_src="manual",
    )
    session.add(account)
    await session.flush()
    mailbox = Mailbox(account_id=account.id, user_id=user.id, path="INBOX", special_use="inbox")
    session.add(mailbox)
    await session.commit()
    try:
        yield Fixture(session, user.id, account.id, mailbox.id)
    finally:
        await session.execute(delete(User).where(User.id == user.id))
        await session.commit()
        await session.close()
        await engine.dispose()


def _msg(
    subject: str,
    body: str,
    *,
    message_id: str,
    frm: str = "Alice <alice@example.com>",
    html: str | None = None,
) -> bytes:
    m = EmailMessage()
    m["Subject"] = subject
    m["From"] = frm
    m["To"] = "me@example.com"
    m["Message-ID"] = message_id
    m["Date"] = "Mon, 09 Mar 2026 08:36:37 +0100"
    m.set_content(body)
    if html:
        m.add_alternative(html, subtype="html")
    return m.as_bytes()


async def test_ingest_creates_rows_and_thread(fx: Fixture) -> None:
    msg = await fx.ingest(
        _msg("Invoice 2026", "Please find the invoice attached.", message_id="<a@example.com>")
    )
    assert msg is not None
    assert msg.thread_id is not None
    body = (await fx.session.scalars(select(BodyPart).where(BodyPart.message_id == msg.id))).all()
    kinds = {b.kind for b in body}
    assert "text" in kinds
    # known_recipients harvested the sender
    kr = await fx.session.scalar(
        select(KnownRecipient).where(KnownRecipient.email == "alice@example.com")
    )
    assert kr is not None and kr.times_used == 1


async def test_idempotent_reingest(fx: Fixture) -> None:
    raw = _msg("Dup", "body", message_id="<dup@example.com>")
    first = await fx.ingest(raw)
    assert first is not None
    # Same (mailbox_id, uid) — but ingest increments uid; force same uid path:
    remote = RemoteMessage(uid=1, raw=raw, flags=frozenset(), internal_date=datetime.now(UTC))
    dup = await persist_message(
        fx.session,
        user_id=fx.user_id,
        account_id=fx.account_id,
        mailbox_id=fx.mailbox_id,
        mailbox_special_use="inbox",
        remote=remote,
    )
    assert dup is None  # (mailbox_id, uid=1) already exists


async def test_search_finds_indexed_message(fx: Fixture) -> None:
    await fx.ingest(
        _msg("Quarterly report", "revenue numbers and projections", message_id="<r@example.com>")
    )
    await fx.ingest(_msg("Lunch plans", "let us grab lunch tomorrow", message_id="<l@example.com>"))

    q = to_search_query("revenue", fx.user_id)
    hits, cursor = await fx.backend.query(q, limit=50, cursor=None)
    assert len(hits) == 1
    assert hits[0].highlights  # revenue highlighted

    q2 = to_search_query("lunch", fx.user_id)
    hits2, _ = await fx.backend.query(q2, limit=50, cursor=None)
    assert len(hits2) == 1


async def test_search_operator_from_filter(fx: Fixture) -> None:
    await fx.ingest(
        _msg("Hi", "hello there", message_id="<h1@example.com>", frm="Bob <bob@example.com>")
    )
    await fx.ingest(
        _msg("Yo", "hello again", message_id="<h2@example.com>", frm="Carol <carol@example.com>")
    )
    q = to_search_query("hello from:bob", fx.user_id)
    hits, _ = await fx.backend.query(q, limit=50, cursor=None)
    assert len(hits) == 1


async def test_html_sanitized_stored(fx: Fixture) -> None:
    msg = await fx.ingest(
        _msg(
            "HTML",
            "text fallback",
            message_id="<html@example.com>",
            html="<p>Hello <script>evil()</script><b>world</b></p>",
        )
    )
    assert msg is not None
    html_part = await fx.session.scalar(
        select(BodyPart).where(BodyPart.message_id == msg.id, BodyPart.kind == "html_sanitized")
    )
    assert html_part is not None
    assert "<script" not in html_part.content.lower()
    assert "world" in html_part.content