86 lines
3.1 KiB
Python
86 lines
3.1 KiB
Python
import pytest
|
|
from datasets import Dataset
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.store.models.document import Document
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rebuild_database(qa_corpus: Dataset, temp_db_path):
|
|
"""Test rebuild functionality with existing documents."""
|
|
async with HaikuRAG(temp_db_path) as client:
|
|
created_docs: list[Document] = []
|
|
for content in qa_corpus["document_extracted"][:3]:
|
|
doc = await client.create_document(
|
|
content=content,
|
|
)
|
|
created_docs.append(doc)
|
|
|
|
documents_before = await client.list_documents()
|
|
assert len(documents_before) == 3
|
|
|
|
chunks_before = []
|
|
for doc in created_docs:
|
|
assert doc.id is not None
|
|
doc_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
chunks_before.extend(doc_chunks)
|
|
|
|
assert len(chunks_before) > 0
|
|
|
|
# Perform rebuild
|
|
processed_doc_ids = []
|
|
async for doc_id in client.rebuild_database():
|
|
processed_doc_ids.append(doc_id)
|
|
|
|
# Verify all documents were processed
|
|
expected_doc_ids = [doc.id for doc in created_docs]
|
|
assert set(processed_doc_ids) == set(expected_doc_ids)
|
|
|
|
documents_after = await client.list_documents()
|
|
assert len(documents_after) == 3
|
|
|
|
# Verify chunks were recreated
|
|
chunks_after = []
|
|
for doc in documents_after:
|
|
if doc.id is not None:
|
|
doc_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
chunks_after.extend(doc_chunks)
|
|
|
|
assert len(chunks_after) > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rebuild_with_missing_source(qa_corpus: Dataset, temp_db_path):
|
|
"""Test rebuild functionality when document source is missing."""
|
|
async with HaikuRAG(temp_db_path) as client:
|
|
# Create document with content
|
|
content = qa_corpus["document_extracted"][0]
|
|
doc = await client.create_document(content=content)
|
|
|
|
# Manually set a URI that doesn't exist
|
|
assert doc.id is not None
|
|
doc_with_uri = await client.document_repository.get_by_id(doc.id)
|
|
assert doc_with_uri is not None
|
|
doc_with_uri.uri = "file:///nonexistent/path.txt"
|
|
await client.document_repository.update(doc_with_uri)
|
|
|
|
# Verify chunks exist before rebuild
|
|
chunks_before = await client.chunk_repository.get_by_document_id(doc.id)
|
|
assert len(chunks_before) > 0
|
|
|
|
# Perform rebuild
|
|
processed_doc_ids = []
|
|
async for doc_id in client.rebuild_database():
|
|
processed_doc_ids.append(doc_id)
|
|
|
|
# Document should still be processed (not skipped)
|
|
assert doc.id in processed_doc_ids
|
|
|
|
# Verify document still exists
|
|
doc_after = await client.document_repository.get_by_id(doc.id)
|
|
assert doc_after is not None
|
|
assert doc_after.content == content
|
|
|
|
# Verify chunks were recreated from content
|
|
chunks_after = await client.chunk_repository.get_by_document_id(doc.id)
|
|
assert len(chunks_after) > 0
|