103 lines
3.0 KiB
Python
103 lines
3.0 KiB
Python
"""Focused tests for document metadata lookup and aggregation."""
|
|
|
|
import pytest
|
|
|
|
from kb import database
|
|
from kb.routes import documents
|
|
|
|
|
|
@pytest.fixture
|
|
def document_db(tmp_path, monkeypatch):
|
|
db_path = tmp_path / "kb.db"
|
|
conn = database.get_connection(db_path)
|
|
database.init_schema(conn, 3)
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO documents(title, source_path, content_hash, doc_type, original_filename)
|
|
VALUES ('Vehicle Guide', '/data/staging/random.pdf', 'hash', 'pdf', 'M38T_manual.pdf')
|
|
"""
|
|
)
|
|
doc_id = conn.execute("SELECT id FROM documents").fetchone()["id"]
|
|
conn.execute(
|
|
"INSERT INTO chunks(document_id, chunk_index, text, enriched_text) VALUES (?, 0, 'body', 'body')",
|
|
(doc_id,),
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
monkeypatch.setattr(documents.cfg, "data_dir", tmp_path)
|
|
return doc_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_filters_title_and_original_filename(document_db):
|
|
by_title = await documents.list_documents(
|
|
type=None, tags=None, title="vehicle", filename=None
|
|
)
|
|
by_filename = await documents.list_documents(
|
|
type=None, tags=None, title=None, filename="m38t"
|
|
)
|
|
|
|
assert [item["id"] for item in by_title] == [document_db]
|
|
assert by_filename[0]["original_filename"] == "M38T_manual.pdf"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_info_can_omit_chunks_without_losing_count(document_db):
|
|
result = await documents.get_document(document_db, include_chunks=False)
|
|
|
|
assert result["chunk_count"] == 1
|
|
assert result["chunks"] == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_find_aggregates_chunk_hits_by_document(monkeypatch):
|
|
class Connection:
|
|
def close(self):
|
|
pass
|
|
|
|
monkeypatch.setattr(documents, "get_connection", lambda _path: Connection())
|
|
monkeypatch.setattr(
|
|
documents,
|
|
"hybrid_search",
|
|
lambda *_args, **_kwargs: {
|
|
"results": [
|
|
{
|
|
"document_id": 7,
|
|
"title": "Guide",
|
|
"doc_type": "pdf",
|
|
"source_path": "/staging/file",
|
|
"original_filename": "guide.pdf",
|
|
"tags": [],
|
|
"score": 0.8,
|
|
"text": "best hit",
|
|
},
|
|
{
|
|
"document_id": 7,
|
|
"title": "Guide",
|
|
"doc_type": "pdf",
|
|
"source_path": "/staging/file",
|
|
"original_filename": "guide.pdf",
|
|
"tags": [],
|
|
"score": 0.7,
|
|
"text": "second hit",
|
|
},
|
|
]
|
|
},
|
|
)
|
|
|
|
result = await documents.find_documents(
|
|
documents.DocumentFindRequest(query="guide")
|
|
)
|
|
|
|
assert result == [{
|
|
"document_id": 7,
|
|
"title": "Guide",
|
|
"doc_type": "pdf",
|
|
"source_path": "/staging/file",
|
|
"original_filename": "guide.pdf",
|
|
"tags": [],
|
|
"score": 0.8,
|
|
"hit_count": 2,
|
|
"top_chunk": "best hit",
|
|
}]
|