dp-zp-agent/tools/rag/tests/test_database.py

from __future__ import annotations

import json
import sqlite3
from pathlib import Path

import pytest

import tools.rag.build_sqlite_index as indexer


def write_json(path: Path, value) -> None:
    path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    path.write_text(
        json.dumps(
            value,
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )


def sample_documents() -> list[dict]:
    return [
        {
            "path": "pages/test/README.md",
            "title": "Strojový preklad",
            "author": "Autor",
            "published": True,
            "content_length": 100,
            "metadata": {
                "published": True,
            },
        }
    ]


def sample_chunks() -> list[dict]:
    return [
        {
            "chunk_id": ("pages/test/README.md::chunk-0"),
            "document_path": ("pages/test/README.md"),
            "title": "Strojový preklad",
            "author": "Autor",
            "published": True,
            "chunk_index": 0,
            "heading_paths": [["Úvod"]],
            "text": ("Dokument: Strojový preklad. " "Neurónový preklad textu."),
            "text_length": 58,
            "token_count": 16,
            "content_hash": "abc",
            "tags": [
                "translation",
                "nlp",
            ],
            "categories": [
                "project",
            ],
        }
    ]


def configure_indexer(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> tuple[
    Path,
    Path,
    Path,
]:
    documents_file = tmp_path / "documents.json"

    chunks_file = tmp_path / "chunks.json"

    db_file = tmp_path / "zp_index.sqlite"

    write_json(
        documents_file,
        sample_documents(),
    )

    write_json(
        chunks_file,
        sample_chunks(),
    )

    monkeypatch.setattr(
        indexer,
        "DOCUMENTS_FILE",
        documents_file,
    )

    monkeypatch.setattr(
        indexer,
        "CHUNKS_FILE",
        chunks_file,
    )

    monkeypatch.setattr(
        indexer,
        "DB_FILE",
        db_file,
    )

    return (
        documents_file,
        chunks_file,
        db_file,
    )


def test_database_contains_documents_chunks_metadata_and_fts(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _, _, db_file = configure_indexer(
        monkeypatch,
        tmp_path,
    )

    counts = indexer.build_database()

    assert counts == {
        "documents": 1,
        "chunks": 1,
        "fts_chunks": 1,
        "embedding_chunks": 1,
        "embedding_segments": 1,
        "embedding_reused": 0,
        "embedding_computed": 1,
        "tags": 2,
        "categories": 1,
    }

    with sqlite3.connect(db_file) as conn:
        assert conn.execute("PRAGMA integrity_check").fetchone()[0] == "ok"

        assert conn.execute("PRAGMA foreign_key_check").fetchall() == []

        assert conn.execute("""
                SELECT published
                FROM chunks
                """).fetchone()[0] == 1

        assert conn.execute("""
                SELECT COUNT(*)
                FROM chunk_tags
                """).fetchone()[0] == 2

        assert conn.execute("""
                SELECT COUNT(*)
                FROM chunk_categories
                """).fetchone()[0] == 1

        assert conn.execute("""
                SELECT COUNT(*)
                FROM chunks_fts
                WHERE chunks_fts
                MATCH 'strojovy'
                """).fetchone()[0] == 1

        assert conn.execute("""
                SELECT COUNT(*)
                FROM chunk_embeddings
                """).fetchone()[0] == 1

        embedding_row = conn.execute("""
            SELECT
                model,
                dimensions,
                LENGTH(embedding)
            FROM chunk_embeddings
            """).fetchone()

        assert embedding_row is not None

        model, dimensions, blob_size = embedding_row

        assert model
        assert dimensions > 0

        # float32 = 4 bajty.
        assert blob_size == dimensions * 4


def test_database_rebuild_replaces_old_database_only_after_success(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _, _, db_file = configure_indexer(
        monkeypatch,
        tmp_path,
    )

    with sqlite3.connect(db_file) as conn:
        conn.execute("""
            CREATE TABLE marker(
                value TEXT
            )
            """)

        conn.execute("""
            INSERT INTO marker
            VALUES ('old database')
            """)

        conn.commit()

    def fail_validation(
        conn: sqlite3.Connection,
    ) -> None:
        raise RuntimeError("úmyselná chyba validácie")

    monkeypatch.setattr(
        indexer,
        "validate_database",
        fail_validation,
    )

    with pytest.raises(
        RuntimeError,
        match="úmyselná chyba",
    ):
        indexer.build_database()

    with sqlite3.connect(db_file) as conn:
        assert conn.execute("""
                SELECT value
                FROM marker
                """).fetchone()[0] == "old database"

    assert not (indexer.temporary_database_path(db_file).exists())


def test_database_rejects_chunk_without_chunk_id(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    (
        documents_file,
        chunks_file,
        _,
    ) = configure_indexer(
        monkeypatch,
        tmp_path,
    )

    broken = sample_chunks()

    broken[0]["chunk_id"] = ""

    write_json(
        documents_file,
        sample_documents(),
    )

    write_json(
        chunks_file,
        broken,
    )

    with pytest.raises(
        ValueError,
        match="chunk_id",
    ):
        indexer.build_database()