dp-zp-agent/test/test_search.py
2026-07-28 23:56:03 +02:00

286 lines
8.0 KiB
Python

from __future__ import annotations
import json
import sqlite3
from pathlib import Path
import pytest
import scripts.build_sqlite_index as indexer
from scripts.search_utils import search_database
def write_json(path: Path, value) -> None:
path.write_text(json.dumps(value, ensure_ascii=False), encoding="utf-8")
def build_search_database(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> Path:
documents_file = tmp_path / "documents.json"
chunks_file = tmp_path / "chunks.json"
db_file = tmp_path / "search.sqlite"
documents = [
{
"path": "pages/students/jan_ptak/README.md",
"title": "Ján Pták",
"author": "Daniel Hladek",
"published": True,
"content_length": 100,
"metadata": {},
},
{
"path": "pages/students/jan_holp/README.md",
"title": "Ján Holp",
"author": "Daniel Hladek",
"published": True,
"content_length": 100,
"metadata": {},
},
{
"path": "pages/topics/translation/README.md",
"title": "Strojový preklad",
"author": "Daniel Hladek",
"published": True,
"content_length": 200,
"metadata": {},
},
{
"path": "pages/topics/open/README.md",
"title": "Otvorené projekty",
"author": "Daniel Hladek",
"published": True,
"content_length": 100,
"metadata": {},
},
{
"path": "pages/topics/deep/README.md",
"title": "Neurónové siete",
"author": "Daniel Hladek",
"published": False,
"content_length": 100,
"metadata": {},
},
]
def chunk(
chunk_id: str,
path: str,
title: str,
text: str,
published: bool,
index: int = 0,
tags: list[str] | None = None,
categories: list[str] | None = None,
) -> dict:
return {
"chunk_id": chunk_id,
"document_path": path,
"title": title,
"author": "Daniel Hladek",
"published": published,
"chunk_index": index,
"heading_paths": [],
"text": text,
"text_length": len(text),
"token_count": 30,
"content_hash": chunk_id,
"tags": tags or [],
"categories": categories or [],
}
chunks = [
chunk(
"jan-ptak::0",
"pages/students/jan_ptak/README.md",
"Ján Pták",
"Dokument: Ján Pták. Agent pre manažment záverečných prác.",
True,
tags=["rag", "nlp"],
categories=["dp2027"],
),
chunk(
"jan-holp::0",
"pages/students/jan_holp/README.md",
"Ján Holp",
"Dokument: Ján Holp. Získavanie informácií a PageRank.",
True,
tags=["ir"],
),
chunk(
"translation::0",
"pages/topics/translation/README.md",
"Strojový preklad",
"Dokument: Strojový preklad. Štatistický strojový preklad.",
True,
index=0,
tags=["translation"],
categories=["project"],
),
chunk(
"translation::1",
"pages/topics/translation/README.md",
"Strojový preklad",
"Dokument: Strojový preklad. Neurónový preklad viet.",
True,
index=1,
tags=["translation"],
categories=["project"],
),
chunk(
"open::0",
"pages/topics/open/README.md",
"Otvorené projekty",
"Téma pre strojový preklad slovenského jazyka.",
True,
categories=["info"],
),
chunk(
"deep::0",
"pages/topics/deep/README.md",
"Neurónové siete",
"Modely hlbokého učenia a trénovanie hlbokých neurónových sietí.",
False,
tags=["nn"],
),
]
write_json(documents_file, documents)
write_json(chunks_file, chunks)
monkeypatch.setattr(indexer, "DOCUMENTS_FILE", documents_file)
monkeypatch.setattr(indexer, "CHUNKS_FILE", chunks_file)
monkeypatch.setattr(indexer, "DB_FILE", db_file)
indexer.build_database()
return db_file
def test_person_query_does_not_add_any_term_noise(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
db_file = build_search_database(tmp_path, monkeypatch)
response = search_database(db_file, "jan ptak", limit=5)
assert response["strategies"] == ["all_terms"]
assert [item["title"] for item in response["results"]] == ["Ján Pták"]
def test_topic_query_returns_best_topic_first(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
db_file = build_search_database(tmp_path, monkeypatch)
response = search_database(db_file, "strojovy preklad", limit=5)
assert response["strategies"] == ["all_terms"]
assert response["results"][0]["title"] == "Strojový preklad"
def test_prefix_fallback_handles_slovak_inflection(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
db_file = build_search_database(tmp_path, monkeypatch)
response = search_database(db_file, "hlboke ucenie", limit=5)
assert response["strategies"] == ["prefix_terms"]
assert response["results"][0]["title"] == "Neurónové siete"
def test_any_term_is_used_only_when_stricter_queries_find_nothing(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
db_file = build_search_database(tmp_path, monkeypatch)
response = search_database(db_file, "nezmysel preklad", limit=5)
assert response["strategies"] == ["any_term"]
assert response["results"]
assert all("preklad" in item["text"].casefold() for item in response["results"])
def test_published_only_excludes_unpublished_chunks(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
db_file = build_search_database(tmp_path, monkeypatch)
all_results = search_database(db_file, "hlboke ucenie", limit=5)
public_results = search_database(
db_file,
"hlboke ucenie",
limit=5,
published_only=True,
)
assert all_results["results"]
assert public_results["results"] == []
def test_results_are_diversified_by_document(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
db_file = build_search_database(tmp_path, monkeypatch)
response = search_database(
db_file,
"preklad",
limit=5,
max_per_document=1,
)
paths = [item["document_path"] for item in response["results"]]
assert len(paths) == len(set(paths))
def test_public_result_format_has_no_internal_id_and_uses_boolean(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
db_file = build_search_database(tmp_path, monkeypatch)
result = search_database(db_file, "jan ptak", limit=1)["results"][0]
assert "id" not in result
assert result["published"] is True
assert result["chunk_id"] == "jan-ptak::0"
assert result["source_url"].endswith("students/jan_ptak")
def test_empty_query_returns_empty_result(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
db_file = build_search_database(tmp_path, monkeypatch)
response = search_database(db_file, " ", limit=5)
assert response == {
"engine": "sqlite_fts5",
"strategies": [],
"results": [],
}
def test_missing_database_is_reported(tmp_path: Path) -> None:
with pytest.raises(FileNotFoundError):
search_database(tmp_path / "missing.sqlite", "rag")
def test_missing_fts_schema_is_reported(tmp_path: Path) -> None:
db_file = tmp_path / "broken.sqlite"
with sqlite3.connect(db_file) as conn:
conn.execute("CREATE TABLE chunks(id INTEGER PRIMARY KEY)")
with pytest.raises(RuntimeError, match="FTS5 index"):
search_database(db_file, "rag")