pablan/backend/tests/test_indexing.py

import pytest
from sqlalchemy import delete, func, select
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker

from app.ingestion.handlers import INDEX_DOCUMENT, REINDEX_ALL
from app.ingestion.queue import enqueue, process_one
from app.models import (
    Chunk,
    Document,
    DocumentStatus,
    DocumentVisibility,
    Job,
    JobStatus,
    User,
)
from app.rag.indexing import reindex_document

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

CONTENT = (
    "## Wartung\n\nWöchentlich schmieren mit GX-220.\n\n"
    "## Sicherheit\n\nLichtvorhang niemals überbrücken.\n"
)


@pytest.fixture
def session_factory(db_engine: AsyncEngine) -> async_sessionmaker[AsyncSession]:
    return async_sessionmaker(db_engine, expire_on_commit=False)


async def _doc(
    db: AsyncSession,
    seeded_user: User,
    *,
    title: str = "Handbuch",
    status: DocumentStatus = DocumentStatus.published,
) -> Document:
    document = Document(
        title=title,
        status=status,
        visibility=DocumentVisibility.department,
        content_md=CONTENT,
        author_id=seeded_user.id,
        department_id=seeded_user.department_id,
        meta={},
    )
    db.add(document)
    await db.flush()
    return document


async def _chunk_count(db: AsyncSession, document_id) -> int:
    return (
        await db.execute(
            select(func.count(Chunk.id)).where(Chunk.document_id == document_id)
        )
    ).scalar_one()


async def test_reindex_creates_chunks_with_meta(
    db: AsyncSession, seeded_user: User
) -> None:
    document = await _doc(db, seeded_user)
    count = await reindex_document(db, document)
    await db.commit()
    assert count == 2

    chunks = (
        (await db.execute(select(Chunk).where(Chunk.document_id == document.id)))
        .scalars()
        .all()
    )
    assert {chunk.meta["heading_path"] for chunk in chunks} == {
        "Handbuch › Wartung",
        "Handbuch › Sicherheit",
    }
    for chunk in chunks:
        assert chunk.meta["visibility"] == "department"
        assert chunk.meta["department_id"] == str(seeded_user.department_id)


async def test_reindex_is_idempotent(db: AsyncSession, seeded_user: User) -> None:
    document = await _doc(db, seeded_user)
    await reindex_document(db, document)
    await db.commit()
    first_ids = {
        chunk.id
        for chunk in (
            await db.execute(select(Chunk).where(Chunk.document_id == document.id))
        ).scalars()
    }

    await reindex_document(db, document)
    await db.commit()
    chunks = (
        (await db.execute(select(Chunk).where(Chunk.document_id == document.id)))
        .scalars()
        .all()
    )
    assert len(chunks) == 2
    assert first_ids.isdisjoint({chunk.id for chunk in chunks})


async def test_index_job_indexes_published_and_cleans_unpublished(
    db: AsyncSession,
    seeded_user: User,
    session_factory: async_sessionmaker[AsyncSession],
) -> None:
    document = await _doc(db, seeded_user)
    await enqueue(db, INDEX_DOCUMENT, {"document_id": str(document.id)})
    await db.commit()
    assert await process_one(session_factory) is True
    assert await _chunk_count(db, document.id) == 2

    document.status = DocumentStatus.archived
    await enqueue(db, INDEX_DOCUMENT, {"document_id": str(document.id)})
    await db.commit()
    assert await process_one(session_factory) is True
    assert await _chunk_count(db, document.id) == 0


async def test_reindex_all_fans_out_and_rebuilds_from_markdown(
    db: AsyncSession,
    seeded_user: User,
    session_factory: async_sessionmaker[AsyncSession],
) -> None:
    published = [await _doc(db, seeded_user, title=f"Doc {i}") for i in range(3)]
    await _doc(db, seeded_user, title="Entwurf", status=DocumentStatus.draft)
    # Simulate an embedding-model swap: all derivatives are gone.
    await db.execute(delete(Chunk))
    await enqueue(db, REINDEX_ALL)
    await db.commit()

    # Fan-out: the reindex_all job only enqueues per-document jobs.
    assert await process_one(session_factory) is True
    index_jobs = (
        (
            await db.execute(
                select(Job).where(
                    Job.type == INDEX_DOCUMENT, Job.status == JobStatus.pending
                )
            )
        )
        .scalars()
        .all()
    )
    assert len(index_jobs) == 3  # the draft is not indexed

    while await process_one(session_factory):
        pass
    for document in published:
        assert await _chunk_count(db, document.id) == 2