webmail-client/backend/tests/search/test_search_api.py

"""Operator search through the API (SPEC §8.2, §9.1), on the Postgres backend."""

from __future__ import annotations

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

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

from app.config import get_settings
from app.identity.passwords import hash_password
from app.identity.sessions import create_session
from app.indexer.indexer import persist_message
from app.main import create_app
from app.models import Account, Mailbox, User
from app.providers.base import RemoteMessage
from app.search.pg import PgSearchBackend


@dataclass
class Ctx:
    client: httpx.AsyncClient
    session: AsyncSession
    user_id: uuid.UUID
    account_id: uuid.UUID
    mailbox_id: uuid.UUID
    uid: int = 0


@pytest_asyncio.fixture
async def ctx() -> AsyncIterator[Ctx]:
    engine = create_async_engine(get_settings().database_url, poolclass=NullPool)
    sm = async_sessionmaker(engine, expire_on_commit=False)
    session = sm()
    user = User(handle=f"sq{uuid.uuid4().hex[:8]}@test.local", password_hash=hash_password("pw"))
    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()
    token = await create_session(user.id, "pytest")
    client = httpx.AsyncClient(
        transport=ASGITransport(app=create_app()),
        base_url="http://testserver",
        cookies={"session": token},
    )
    try:
        yield Ctx(
            client=client,
            session=session,
            user_id=user.id,
            account_id=account.id,
            mailbox_id=mailbox.id,
        )
    finally:
        await client.aclose()
        await session.execute(delete(User).where(User.id == user.id))
        await session.commit()
        await session.close()
        await engine.dispose()


async def _ingest(
    ctx: Ctx,
    *,
    frm: str,
    subject: str,
    body: str,
    flags: frozenset[str] = frozenset(),
    attach: bool = False,
) -> None:
    ctx.uid += 1
    m = EmailMessage()
    m["Subject"] = subject
    m["From"] = frm
    m["To"] = "me@example.com"
    m["Message-ID"] = f"<{uuid.uuid4().hex}@x.com>"
    m["Date"] = "Mon, 09 Mar 2026 08:36:37 +0100"
    m.set_content(body)
    if attach:
        m.add_attachment(b"data", maintype="application", subtype="pdf", filename="report.pdf")
    result = await persist_message(
        ctx.session,
        user_id=ctx.user_id,
        account_id=ctx.account_id,
        mailbox_id=ctx.mailbox_id,
        mailbox_special_use="inbox",
        remote=RemoteMessage(
            uid=ctx.uid, raw=m.as_bytes(), flags=flags, internal_date=datetime.now(UTC)
        ),
    )
    await ctx.session.commit()
    assert result is not None
    await PgSearchBackend().index(result.search_doc)


async def _search(ctx: Ctx, q: str) -> int:
    resp = await ctx.client.get("/api/search", params={"q": q})
    assert resp.status_code == 200, resp.text
    return len(resp.json()["threads"])


@pytest_asyncio.fixture
async def seeded(ctx: Ctx) -> Ctx:
    await _ingest(
        ctx,
        frm="Bob <bob@corp.com>",
        subject="Quarterly revenue",
        body="revenue numbers up",
        flags=frozenset(),
    )  # unread
    await _ingest(
        ctx,
        frm="Carol <carol@corp.com>",
        subject="Lunch",
        body="grab lunch",
        flags=frozenset({"\\Seen"}),
    )  # read
    await _ingest(
        ctx,
        frm="Dave <dave@corp.com>",
        subject="Report attached",
        body="see the pdf",
        flags=frozenset({"\\Seen"}),
        attach=True,
    )
    return ctx


async def test_free_text(seeded: Ctx) -> None:
    assert await _search(seeded, "revenue") == 1
    assert await _search(seeded, "lunch") == 1


async def test_free_text_prefix_matches_while_typing(seeded: Ctx) -> None:
    # Search-as-you-type: a partial word finds the message before it is complete.
    assert await _search(seeded, "reven") == 1  # -> "revenue"
    assert await _search(seeded, "lun") == 1  # -> "lunch"
    # Prefix is anchored at word start, not a substring match.
    assert await _search(seeded, "enue") == 0
    # Multiple partial words are ANDed as prefixes.
    assert await _search(seeded, "quart reven") == 1  # "Quarterly revenue"
    assert await _search(seeded, "quart lun") == 0


async def test_free_text_full_email_address(seeded: Ctx) -> None:
    # Postgres keeps an email as a single lexeme; searching the *whole* address
    # must still match (regression: prefix-splitting on word chars broke this).
    assert await _search(seeded, "bob@corp.com") == 1
    assert await _search(seeded, "bob@corp.co") == 1  # prefix while still typing


async def test_free_text_email_matches_address_consistently(seeded: Ctx) -> None:
    # A term with '@' is treated as an address search and matches as a substring
    # of from/to/cc, so it narrows *predictably* as you type more of the address
    # (a partial term never returns more than the fuller one).
    assert await _search(seeded, "bob@") == 1  # local part -> bob only
    assert await _search(seeded, "bob@corp") == 1
    assert await _search(seeded, "bob@co") == 1
    assert await _search(seeded, "@corp.com") == 3  # all three share the domain
    assert await _search(seeded, "zzz@corp") == 0  # no such sender


async def test_address_matches_rank_before_body_matches(ctx: Ctx) -> None:
    # Older message matches "zebra" only via its sender address; the newer one
    # matches only in the body. The address/name hit must still come first.
    await _ingest(ctx, frm="Zebra Corp <hello@zebra.com>", subject="hi", body="nothing here")
    await _ingest(ctx, frm="Someone <someone@corp.com>", subject="a note", body="the zebra ran")
    resp = await ctx.client.get("/api/search", params={"q": "zebra"})
    subjects = [t["thread"]["subject"] for t in resp.json()["threads"]]
    assert subjects == ["hi", "a note"]  # address hit first, then body-only


async def test_free_text_bare_domain_word_matches_address(seeded: Ctx) -> None:
    # A bare domain/local word finds the address even though Postgres indexes the
    # whole address as one lexeme — address parts are tokenized at index time.
    assert await _search(seeded, "corp") == 3  # the domain word -> all three
    assert await _search(seeded, "carol") == 1  # local part -> carol only
    # A dotted domain term also matches via the address-substring fallback.
    assert await _search(seeded, "corp.com") == 3


async def test_from_operator(seeded: Ctx) -> None:
    assert await _search(seeded, "from:bob") == 1
    assert await _search(seeded, "from:corp.com") == 3  # all three share the domain


async def test_is_unread_operator(seeded: Ctx) -> None:
    assert await _search(seeded, "is:unread") == 1
    assert await _search(seeded, "is:read") == 2


async def test_has_attachment_operator(seeded: Ctx) -> None:
    assert await _search(seeded, "has:attachment") == 1


async def test_subject_operator(seeded: Ctx) -> None:
    assert await _search(seeded, "subject:revenue") == 1


async def test_negation(seeded: Ctx) -> None:
    # everything unread that is NOT from bob -> none
    assert await _search(seeded, "is:unread -from:bob") == 0


async def test_combined_operators(seeded: Ctx) -> None:
    assert await _search(seeded, "from:dave has:attachment") == 1
    assert await _search(seeded, "from:bob has:attachment") == 0