pablan/backend/tests/test_retrieval.py

"""Permission boundaries and hybrid plumbing of rag.retrieval.search.

Uses deterministic fake embeddings (fake_embed fixture): identical text →
distance 0; retrieval semantics with real embeddings live in tests/evals.
"""

import uuid

import pytest
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.auth.passwords import hash_password
from app.metrics import metrics
from app.models import (
    Chunk,
    Department,
    DocPermission,
    Document,
    DocumentStatus,
    DocumentVisibility,
    User,
    UserRole,
)
from app.rag.indexing import embedding_text, reindex_document
from app.rag.retrieval import search, text_search

pytestmark = pytest.mark.usefixtures("fake_embed")


class Setup:
    pablo: User  # Engineering
    max: User  # Sales
    norbert: User  # no department
    pub: Document
    dept_eng: Document
    restricted_ben: Document
    granted_eng: Document
    draft: Document


async def _user(db: AsyncSession, email: str, department_id) -> User:
    user = User(
        email=email,
        name=email.split("@")[0],
        role=UserRole.member,
        password_hash=hash_password("secret123"),
        department_id=department_id,
    )
    db.add(user)
    await db.flush()
    return user


async def _doc(
    db: AsyncSession,
    *,
    title: str,
    content: str,
    author: User,
    department_id=None,
    visibility: DocumentVisibility,
    status: DocumentStatus = DocumentStatus.published,
) -> Document:
    document = Document(
        title=title,
        status=status,
        visibility=visibility,
        content_md=content,
        author_id=author.id,
        department_id=department_id,
    )
    db.add(document)
    await db.flush()
    return document


@pytest.fixture
async def setup(db: AsyncSession) -> Setup:
    s = Setup()
    engineering = Department(name="Engineering")
    sales = Department(name="Sales")
    db.add_all([engineering, sales])
    await db.flush()

    s.pablo = await _user(db, "pablo@test.dev", engineering.id)
    s.max = await _user(db, "max@test.dev", sales.id)
    s.norbert = await _user(db, "norbert@test.dev", None)

    s.pub = await _doc(
        db,
        title="Kaffeemaschine",
        content="## Pflege\n\nDie Kaffeemaschine wird jeden Freitag entkalkt.",
        author=s.pablo,
        department_id=engineering.id,
        visibility=DocumentVisibility.public,
    )
    s.dept_eng = await _doc(
        db,
        title="Bandschleifer BS-100",
        content="## Wartung\n\nDer Bandschleifer braucht wöchentlich ein neues Schleifband.",
        author=s.pablo,
        department_id=engineering.id,
        visibility=DocumentVisibility.department,
    )
    s.restricted_ben = await _doc(
        db,
        title="Geheime Preisliste",
        content="## Preise\n\nDer Rabattdeckel liegt bei zwölf Prozent.",
        author=s.max,
        department_id=sales.id,
        visibility=DocumentVisibility.restricted,
    )
    s.granted_eng = await _doc(
        db,
        title="Ersatzteillager",
        content="## Zugang\n\nDie Zugangskarte für das Ersatzteillager liegt im Tresorfach drei.",
        author=s.max,
        department_id=sales.id,
        visibility=DocumentVisibility.restricted,
    )
    db.add(DocPermission(document_id=s.granted_eng.id, department_id=engineering.id))
    s.draft = await _doc(
        db,
        title="Pausenregelung",
        content="## Entwurf\n\nNeue Pausenregelung ab Oktober.",
        author=s.pablo,
        department_id=engineering.id,
        visibility=DocumentVisibility.public,
        status=DocumentStatus.draft,
    )

    for document in (s.pub, s.dept_eng, s.restricted_ben, s.granted_eng, s.draft):
        await reindex_document(db, document)
    await db.commit()
    return s


def _doc_ids(results) -> set[uuid.UUID]:
    return {result.document_id for result in results}


async def test_results_carry_citation_metadata(db: AsyncSession, setup: Setup) -> None:
    # Every query term must exist in the target chunk: websearch_to_tsquery
    # ANDs terms, and the fake embeddings carry no semantics.
    results = await search(db, "Kaffeemaschine entkalkt", user=setup.pablo)
    assert results, "public document not found"
    top = results[0]
    assert top.document_id == setup.pub.id
    assert top.title == "Kaffeemaschine"
    assert top.heading_path == "Kaffeemaschine › Pflege"
    assert top.content
    assert top.score > 0


async def test_department_visibility(db: AsyncSession, setup: Setup) -> None:
    query = "Schleifband für den Bandschleifer"
    assert setup.dept_eng.id in _doc_ids(await search(db, query, user=setup.pablo))
    assert setup.dept_eng.id not in _doc_ids(await search(db, query, user=setup.max))
    assert setup.dept_eng.id not in _doc_ids(
        await search(db, query, user=setup.norbert)
    )


async def test_restricted_needs_grant_or_authorship(
    db: AsyncSession, setup: Setup
) -> None:
    query = "Zugangskarte Ersatzteillager Tresorfach"
    # Engineering has an explicit grant.
    assert setup.granted_eng.id in _doc_ids(await search(db, query, user=setup.pablo))
    # The author always sees their own document.
    assert setup.granted_eng.id in _doc_ids(await search(db, query, user=setup.max))
    # No department, no grant, no authorship: nothing.
    assert setup.granted_eng.id not in _doc_ids(
        await search(db, query, user=setup.norbert)
    )


async def test_restricted_document_never_leaks(db: AsyncSession, setup: Setup) -> None:
    """Acceptance: user A can NEVER retrieve chunks of user B's restricted
    document — tested through both retrieval branches."""
    # Full-text branch: the exact distinctive term.
    fts_results = await search(db, "Rabattdeckel", user=setup.pablo)
    assert setup.restricted_ben.id not in _doc_ids(fts_results)

    # Vector branch: query IS the exact chunk content (distance 0 — it would
    # be the top hit if the filter leaked).
    chunk = (
        await db.execute(
            select(Chunk.content, Chunk.meta).where(
                Chunk.document_id == setup.restricted_ben.id
            )
        )
    ).one()
    chunk_content = embedding_text(chunk.meta["heading_path"], chunk.content)
    vec_results = await search(db, chunk_content, user=setup.pablo)
    assert setup.restricted_ben.id not in _doc_ids(vec_results)

    # The author, of course, finds it.
    assert setup.restricted_ben.id in _doc_ids(
        await search(db, "Rabattdeckel", user=setup.max)
    )


async def test_unpublished_documents_are_never_searchable(
    db: AsyncSession, setup: Setup
) -> None:
    """Draft chunks exist in the table but must never surface — not even for
    the author, not even for a query that is the exact chunk content."""
    chunk = (
        await db.execute(
            select(Chunk.content, Chunk.meta).where(Chunk.document_id == setup.draft.id)
        )
    ).one()
    chunk_content = embedding_text(chunk.meta["heading_path"], chunk.content)
    for query in ("Pausenregelung Entwurf", chunk_content):
        assert setup.draft.id not in _doc_ids(await search(db, query, user=setup.pablo))


async def test_status_change_applies_without_reindex(
    db: AsyncSession, setup: Setup
) -> None:
    query = "Zugangskarte Ersatzteillager Tresorfach"
    assert setup.granted_eng.id in _doc_ids(await search(db, query, user=setup.pablo))
    setup.granted_eng.status = DocumentStatus.archived
    await db.commit()
    # Chunks still exist, but the live status filter hides them instantly.
    assert setup.granted_eng.id not in _doc_ids(
        await search(db, query, user=setup.pablo)
    )


async def test_visibility_change_applies_without_reindex(
    db: AsyncSession, setup: Setup
) -> None:
    """The permission filter reads the documents table, never the stale
    denormalized copy in chunk meta."""
    query = "Schleifband für den Bandschleifer"
    assert setup.dept_eng.id in _doc_ids(await search(db, query, user=setup.pablo))
    setup.dept_eng.visibility = DocumentVisibility.restricted
    setup.dept_eng.author_id = setup.max.id  # take authorship out of the way
    await db.commit()
    assert setup.dept_eng.id not in _doc_ids(await search(db, query, user=setup.pablo))


async def test_vector_branch_finds_exact_content(
    db: AsyncSession, setup: Setup
) -> None:
    chunk = (
        await db.execute(
            select(Chunk.content, Chunk.meta).where(Chunk.document_id == setup.pub.id)
        )
    ).one()
    chunk_content = embedding_text(chunk.meta["heading_path"], chunk.content)
    results = await search(db, chunk_content, user=setup.norbert)
    assert results
    top = results[0]
    assert top.document_id == setup.pub.id
    assert top.vector_distance is not None
    assert top.vector_distance < 0.001


async def test_top_k_limits_results(db: AsyncSession, setup: Setup) -> None:
    results = await search(db, "Kaffeemaschine entkalkt", user=setup.pablo, top_k=1)
    assert len(results) <= 1


async def test_search_records_metrics(db: AsyncSession, setup: Setup) -> None:
    await search(db, "Kaffeemaschine", user=setup.pablo)
    snapshot = metrics.snapshot()
    assert snapshot["counters"]["retrieval_searches_total"][0]["value"] >= 1
    assert "retrieval_seconds" in snapshot["histograms"]


async def test_text_search_needs_no_embedding_and_keeps_the_permission_filter(
    db: AsyncSession, setup: Setup, monkeypatch: pytest.MonkeyPatch
) -> None:
    """The fallback for a dead embedding endpoint: keyword matching over the
    tsvector index alone, with the same permission CTE as `search`."""

    async def _no_endpoint(texts: list[str], *, role: str = "embedding"):
        raise AssertionError("text_search must not embed")

    monkeypatch.setattr("app.rag.retrieval.embed", _no_endpoint)

    results = await text_search(db, "Kaffeemaschine entkalkt", user=setup.pablo)
    assert setup.pub.id in _doc_ids(results)
    assert all(result.fts_match for result in results)
    assert all(result.vector_distance is None for result in results)

    # Same boundaries as the hybrid path: someone else's restricted document
    # stays invisible, an unpublished draft stays out.
    assert setup.restricted_ben.id not in _doc_ids(
        await text_search(db, "Rabattdeckel", user=setup.pablo)
    )
    assert setup.restricted_ben.id in _doc_ids(
        await text_search(db, "Rabattdeckel", user=setup.max)
    )


async def test_text_search_returns_nothing_for_an_unmatched_query(
    db: AsyncSession, setup: Setup
) -> None:
    """No fuzzy rescue without vectors: a word nobody wrote finds nothing,
    which is what the UI has to be able to say."""
    assert await text_search(db, "Quantenverschraenkung", user=setup.pablo) == []


async def test_text_search_answers_a_whole_question(
    db: AsyncSession, setup: Setup
) -> None:
    """A question is typed as a sentence, and without a vector half to carry
    the recall, requiring every word in one chunk would find nothing."""
    results = await text_search(
        db, "Wie wird die Kaffeemaschine eigentlich entkalkt?", user=setup.pablo
    )
    assert setup.pub.id in _doc_ids(results)


async def test_text_search_ignores_a_query_of_only_stop_words(
    db: AsyncSession, setup: Setup
) -> None:
    """Nothing to search for is an empty result, not a database error."""
    assert await text_search(db, "und der die", user=setup.pablo) == []