danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Do not use private methods in tests

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-11-25 14:19:28 +02:00
parent c22a5bb0ce
commit 2c554a9350
No known key found for this signature in database
4 changed files with 57 additions and 87 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
tests/test_chunk.py
View file

@ -1,6 +1,7 @@
import pytest
from datasets import Dataset
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.converters import get_converter
from haiku.rag.store.engine import Store
@ -13,41 +14,42 @@ from haiku.rag.store.repositories.document import DocumentRepository
@pytest.mark.asyncio
async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path):
"""Test ChunkRepository operations."""
# Create a store and repositories
store = Store(temp_db_path)
doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store)
# Create client
client = HaikuRAG(db_path=temp_db_path, config=Config)
# Get the first document from the corpus
first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"]
# Create a document first with chunks
document = Document(content=document_text, metadata={"source": "test"})
converter = get_converter(Config)
docling_document = converter.convert_text(document_text, name="test.md")
created_document = await doc_repo._create_and_chunk(document, docling_document)
created_document = await client.create_document(
content=document_text, metadata={"source": "test"}
)
assert created_document.id is not None
# Test getting chunks by document ID
chunks = await chunk_repo.get_by_document_id(created_document.id)
chunks = await client.chunk_repository.get_by_document_id(created_document.id)
assert len(chunks) > 0
assert all(chunk.document_id == created_document.id for chunk in chunks)
# Test chunk search
results = await chunk_repo.search("election", limit=2, search_type="vector")
results = await client.chunk_repository.search(
"election", limit=2, search_type="vector"
)
assert len(results) <= 2
assert all(hasattr(chunk, "content") for chunk, _ in results)
# Test deleting chunks by document ID
deleted = await chunk_repo.delete_by_document_id(created_document.id)
deleted = await client.chunk_repository.delete_by_document_id(created_document.id)
assert deleted is True
# Verify chunks are gone
chunks_after_delete = await chunk_repo.get_by_document_id(created_document.id)
chunks_after_delete = await client.chunk_repository.get_by_document_id(
created_document.id
)
assert len(chunks_after_delete) == 0
store.close()
client.close()
@pytest.mark.asyncio

25
tests/test_document.py
View file

@ -1,8 +1,8 @@
import pytest
from datasets import Dataset
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.converters import get_converter
from haiku.rag.store.engine import Store
from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.document import DocumentRepository
@ -11,36 +11,25 @@ from haiku.rag.store.repositories.document import DocumentRepository
@pytest.mark.asyncio
async def test_create_document_with_chunks(qa_corpus: Dataset, temp_db_path):
"""Test creating a document with chunks from the qa_corpus using repository."""
# Create a store and repository
store = Store(temp_db_path)
doc_repo = DocumentRepository(store)
# Create client
client = HaikuRAG(db_path=temp_db_path, config=Config)
# Get the first document from the corpus
first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"]
# Create a Document instance
document = Document(
# Create the document with chunks in the database
created_document = await client.create_document(
content=document_text,
metadata={"source": "qa_corpus", "topic": first_doc.get("document_topic", "")},
)
# Convert text to DoclingDocument for chunk creation
converter = get_converter(Config)
docling_document = converter.convert_text(document_text, name="test.md")
# Create the document with chunks in the database
created_document = await doc_repo._create_and_chunk(document, docling_document)
# Verify the document was created
assert created_document.id is not None
assert created_document.content == document_text
# Check that chunks were created using repository
from haiku.rag.store.repositories.chunk import ChunkRepository
chunk_repo = ChunkRepository(store)
chunks = await chunk_repo.get_by_document_id(created_document.id)
chunks = await client.chunk_repository.get_by_document_id(created_document.id)
assert len(chunks) > 0
@ -48,7 +37,7 @@ async def test_create_document_with_chunks(qa_corpus: Dataset, temp_db_path):
for i, chunk in enumerate(chunks):
assert chunk.order == i
store.close()
client.close()
@pytest.mark.asyncio

86
tests/test_search.py
View file

@ -1,21 +1,15 @@
import pytest
from datasets import Dataset
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.converters import get_converter
from haiku.rag.store.engine import Store
from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository
@pytest.mark.asyncio
async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
"""Test that documents can be found by searching with their associated questions."""
# Create a store and repositories
store = Store(temp_db_path)
doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store)
# Create client
client = HaikuRAG(db_path=temp_db_path, config=Config)
# Load unique documents (limited to 10)
seen_documents = set()
@ -31,13 +25,8 @@ async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
continue
seen_documents.add(document_id)
# Create a Document instance
document = Document(content=document_text)
# Create the document with chunks and embeddings
converter = get_converter(Config)
docling_document = converter.convert_text(document_text, name="test.md")
created_document = await doc_repo._create_and_chunk(document, docling_document)
created_document = await client.create_document(content=document_text)
documents.append((created_document, doc_data))
# Test with first few unique documents
@ -46,48 +35,45 @@ async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
question = doc_data["question"]
# Test vector search
vector_results = await chunk_repo.search(
vector_results = await client.chunk_repository.search(
question, limit=5, search_type="vector"
)
target_document_ids = {chunk.document_id for chunk, _ in vector_results}
assert target_document.id in target_document_ids
# Test FTS search
fts_results = await chunk_repo.search(question, limit=5, search_type="fts")
fts_results = await client.chunk_repository.search(
question, limit=5, search_type="fts"
)
target_document_ids = {chunk.document_id for chunk, _ in fts_results}
assert target_document.id in target_document_ids
# Test hybrid search
hybrid_results = await chunk_repo.search(
hybrid_results = await client.chunk_repository.search(
question, limit=5, search_type="hybrid"
)
target_document_ids = {chunk.document_id for chunk, _ in hybrid_results}
assert target_document.id in target_document_ids
store.close()
client.close()
@pytest.mark.asyncio
async def test_chunks_include_document_info(temp_db_path):
"""Test that search results include document URI and metadata."""
store = Store(temp_db_path)
doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store)
client = HaikuRAG(db_path=temp_db_path, config=Config)
# Create a document with URI and metadata
document = Document(
created_document = await client.create_document(
content="This is a test document with some content for searching.",
uri="https://example.com/test.html",
metadata={"title": "Test Document", "author": "Test Author"},
)
# Create the document with chunks
converter = get_converter(Config)
docling_document = converter.convert_text(document.content, name="test.md")
created_document = await doc_repo._create_and_chunk(document, docling_document)
# Search for chunks
results = await chunk_repo.search("test document", limit=1, search_type="hybrid")
results = await client.chunk_repository.search(
"test document", limit=1, search_type="hybrid"
)
assert len(results) > 0
chunk, score = results[0]
@ -101,30 +87,25 @@ async def test_chunks_include_document_info(temp_db_path):
assert chunk.document_meta == {"title": "Test Document", "author": "Test Author"}
assert chunk.document_id == created_document.id
store.close()
client.close()
@pytest.mark.asyncio
async def test_chunks_include_document_title(temp_db_path):
"""Test that search results include the parent document title when present."""
store = Store(temp_db_path)
doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store)
client = HaikuRAG(db_path=temp_db_path, config=Config)
# Create a document with URI and title
document = Document(
await client.create_document(
content="This is a test document with a custom title to verify enrichment.",
uri="file:///tmp/title-test.md",
title="My Custom Title",
)
# Create the document with chunks
converter = get_converter(Config)
dl = converter.convert_text(document.content, name="title-test.md")
await doc_repo._create_and_chunk(document, dl)
# Perform a search that should find this document
results = await chunk_repo.search("custom title", limit=3, search_type="hybrid")
results = await client.chunk_repository.search(
"custom title", limit=3, search_type="hybrid"
)
assert results, "Expected at least one search result"
for chunk, _ in results:
@ -132,15 +113,13 @@ async def test_chunks_include_document_title(temp_db_path):
if chunk.document_uri == "file:///tmp/title-test.md":
assert chunk.document_title == "My Custom Title"
store.close()
client.close()
@pytest.mark.asyncio
async def test_search_score_types(temp_db_path):
"""Test that different search types return appropriate score ranges."""
store = Store(temp_db_path)
doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store)
client = HaikuRAG(db_path=temp_db_path, config=Config)
# Create multiple documents with different content
documents_content = [
@ -150,26 +129,29 @@ async def test_search_score_types(temp_db_path):
"Computer vision systems can interpret and analyze visual information from images.",
]
converter = get_converter(Config)
for content in documents_content:
document = Document(content=content)
docling_document = converter.convert_text(content, name="test.md")
await doc_repo._create_and_chunk(document, docling_document)
await client.create_document(content=content)
query = "machine learning"
# Test vector search scores (should be converted from distances)
vector_results = await chunk_repo.search(query, limit=3, search_type="vector")
vector_results = await client.chunk_repository.search(
query, limit=3, search_type="vector"
)
assert len(vector_results) > 0
vector_scores = [score for _, score in vector_results]
# Test FTS search scores (should be native LanceDB FTS scores)
fts_results = await chunk_repo.search(query, limit=3, search_type="fts")
fts_results = await client.chunk_repository.search(
query, limit=3, search_type="fts"
)
assert len(fts_results) > 0
fts_scores = [score for _, score in fts_results]
# Test hybrid search scores (should be native LanceDB relevance scores)
hybrid_results = await chunk_repo.search(query, limit=3, search_type="hybrid")
hybrid_results = await client.chunk_repository.search(
query, limit=3, search_type="hybrid"
)
assert len(hybrid_results) > 0
hybrid_scores = [score for _, score in hybrid_results]
@ -201,4 +183,4 @@ async def test_search_score_types(temp_db_path):
f"{search_type} results should be sorted by score descending"
)
store.close()
client.close()

5
tests/test_versioning.py
View file

@ -212,11 +212,8 @@ async def test_vacuum_completes_before_context_exit(temp_db_path, monkeypatch):
async with HaikuRAG(db_path=temp_db_path) as client:
# Create multiple documents - each creation triggers automatic vacuum with retention=0
# This aggressively cleans up old versions between operations
converter = get_converter(Config)
for i in range(3):
doc = Document(content=f"Test document {i}")
dl_doc = converter.convert_text(f"Test document {i}", name=f"test{i}.md")
await client.document_repository._create_and_chunk(doc, dl_doc)
await client.create_document(content=f"Test document {i}")
# After context exit, automatic vacuum should have kept versions minimal
store = Store(temp_db_path)