49 lines
1.6 KiB
Python
49 lines
1.6 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
|