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,
    mailbox_id: uuid.UUID | None = None,
    special: str = "inbox",
) -> 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=mailbox_id or ctx.mailbox_id,
        mailbox_special_use=special,
        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_search_excludes_trash_by_default(ctx: Ctx) -> None:
    # Deleting a message from search results (it moves to Trash) removes it from the
    # results — search excludes Trash/Junk unless in:trash/in:junk asks for them.
    from app.models import Mailbox

    await _ingest(ctx, frm="Bob <bob@corp.com>", subject="quarterly zephyr", body="hi")
    trash = Mailbox(
        account_id=ctx.account_id, user_id=ctx.user_id, path="Trash", special_use="trash"
    )
    ctx.session.add(trash)
    await ctx.session.flush()
    await _ingest(
        ctx, frm="X <x@corp.com>", subject="zephyr in trash", body="bye",
        mailbox_id=trash.id, special="trash",
    )
    assert await _search(ctx, "zephyr") == 1  # inbox only; the trashed one is hidden
    assert await _search(ctx, "in:trash zephyr") == 1  # explicit -> searches Trash


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


async def test_more_than_one_page_counts_and_pages(ctx: Ctx) -> None:
    """Beyond the first page: the first page says how many conversations match
    in all, the cursor brings the rest, and Unread/Pinned filter conversations."""
    for i in range(7):
        await _ingest(
            ctx,
            frm=f"Sender {i} <s{i}@example.org>",
            subject=f"Quarterly report {i}",
            body="numbers inside",
            flags=frozenset({"\\Seen"}) if i % 2 else frozenset(),
        )
    await _ingest(ctx, frm="Other <o@example.org>", subject="Lunch", body="nothing")

    first = (await ctx.client.get("/api/search", params={"q": "quarterly", "limit": 3})).json()
    assert first["total"] == 7 and len(first["threads"]) == 3 and first["cursor"]
    seen = {t["thread"]["id"] for t in first["threads"]}
    cursor = first["cursor"]
    while cursor:
        page = (
            await ctx.client.get(
                "/api/search", params={"q": "quarterly", "limit": 3, "cursor": cursor}
            )
        ).json()
        assert page["total"] is None  # counted once, on the first page
        seen |= {t["thread"]["id"] for t in page["threads"]}
        cursor = page["cursor"]
    assert len(seen) == 7

    unread = (
        await ctx.client.get("/api/search", params={"q": "quarterly", "filter": "unread"})
    ).json()
    assert unread["total"] == 4 and len(unread["threads"]) == 4
    assert all(t["thread"]["unread_count"] > 0 for t in unread["threads"])
    pinned = (
        await ctx.client.get("/api/search", params={"q": "quarterly", "filter": "pinned"})
    ).json()
    assert pinned["total"] == 0 and pinned["threads"] == []
    bad = await ctx.client.get("/api/search", params={"q": "quarterly", "filter": "nope"})
    assert bad.status_code == 422