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

Improve test coverage

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-24 13:43:51 +03:00
parent d99e48a9a3
commit 258ffa41fa
No known key found for this signature in database
5 changed files with 467 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
haiku_rag_slim/haiku/rag/client/rebuild.py
View file

@ -16,6 +16,8 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
_REBUILD_BATCH_SIZE = 50
async def rebuild_database( async def rebuild_database(
client: "HaikuRAG", mode: "RebuildMode | None" = None client: "HaikuRAG", mode: "RebuildMode | None" = None
@ -204,7 +206,6 @@ async def _rebuild_rechunk(
"""Re-chunk and re-embed from existing document content.""" """Re-chunk and re-embed from existing document content."""
from haiku.rag.embeddings import embed_chunks from haiku.rag.embeddings import embed_chunks
batch_size = 50
pending_chunks: list[Chunk] = [] pending_chunks: list[Chunk] = []
pending_docs: list[Document] = [] pending_docs: list[Document] = []
pending_doc_ids: list[str] = [] pending_doc_ids: list[str] = []
@ -234,7 +235,7 @@ async def _rebuild_rechunk(
pending_doc_ids.append(doc.id) pending_doc_ids.append(doc.id)
# Flush batch when size reached # Flush batch when size reached
if len(pending_docs) >= batch_size: if len(pending_docs) >= _REBUILD_BATCH_SIZE:
await _flush_rebuild_batch(client, pending_docs, pending_chunks) await _flush_rebuild_batch(client, pending_docs, pending_chunks)
for doc_id in pending_doc_ids: for doc_id in pending_doc_ids:
yield doc_id yield doc_id
@ -255,7 +256,6 @@ async def _rebuild_full(
"""Full rebuild: re-convert from source, re-chunk, re-embed.""" """Full rebuild: re-convert from source, re-chunk, re-embed."""
from haiku.rag.embeddings import embed_chunks from haiku.rag.embeddings import embed_chunks
batch_size = 50
pending_chunks: list[Chunk] = [] pending_chunks: list[Chunk] = []
pending_docs: list[Document] = [] pending_docs: list[Document] = []
pending_doc_ids: list[str] = [] pending_doc_ids: list[str] = []
@ -312,7 +312,7 @@ async def _rebuild_full(
pending_doc_ids.append(doc.id) pending_doc_ids.append(doc.id)
# Flush batch when size reached # Flush batch when size reached
if len(pending_docs) >= batch_size: if len(pending_docs) >= _REBUILD_BATCH_SIZE:
await _flush_rebuild_batch(client, pending_docs, pending_chunks) await _flush_rebuild_batch(client, pending_docs, pending_chunks)
for doc_id in pending_doc_ids: for doc_id in pending_doc_ids:
yield doc_id yield doc_id

242
tests/cassettes/test_rebuild/test_rebuild_batch_size_flush.yaml Normal file
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_rebuild/test_rebuild_full_with_accessible_source.yaml Normal file
View file

File diff suppressed because one or more lines are too long

33
tests/test_client.py
View file

@ -1272,6 +1272,39 @@ async def test_client_convert_file_not_found(temp_db_path):
await client.convert(Path("/nonexistent/path/file.txt")) await client.convert(Path("/nonexistent/path/file.txt"))
async def test_client_convert_from_url(temp_db_path):
"""convert() with an http(s) URL downloads to a tempfile and converts."""
from docling_core.types.doc.document import DoclingDocument
async with HaikuRAG(temp_db_path, create=True) as client:
mock_response = AsyncMock()
mock_response.content = (
b"<html><body><p>URL convert path content.</p></body></html>"
)
mock_response.headers = {"content-type": "text/html"}
mock_response.raise_for_status = AsyncMock()
with patch("httpx.AsyncClient.get", return_value=mock_response):
docling_doc = await client.convert("https://example.com/page.html")
assert isinstance(docling_doc, DoclingDocument)
markdown = docling_doc.export_to_markdown()
assert "URL convert path content" in markdown
async def test_client_convert_from_url_unsupported_content_type(temp_db_path):
"""convert() rejects URLs whose content type isn't supported by the converter."""
async with HaikuRAG(temp_db_path, create=True) as client:
mock_response = AsyncMock()
mock_response.content = b"\x00\x01\x02binary"
mock_response.headers = {"content-type": "application/octet-stream"}
mock_response.raise_for_status = AsyncMock()
with patch("httpx.AsyncClient.get", return_value=mock_response):
with pytest.raises(ValueError, match="Unsupported content type"):
await client.convert("https://example.com/blob.bin")
@pytest.mark.vcr() @pytest.mark.vcr()
async def test_client_convert_unsupported_extension(temp_db_path): async def test_client_convert_unsupported_extension(temp_db_path):
"""Test convert() raises ValueError for unsupported file extension.""" """Test convert() raises ValueError for unsupported file extension."""

106
tests/test_rebuild.py
View file

@ -1,3 +1,5 @@
import tempfile
from pathlib import Path
from typing import TypedDict from typing import TypedDict
import pytest import pytest
@ -263,3 +265,107 @@ async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path):
# Chunk IDs should change (chunks are recreated) # Chunk IDs should change (chunks are recreated)
assert chunk_ids_before.isdisjoint(chunk_ids_after) assert chunk_ids_before.isdisjoint(chunk_ids_after)
@pytest.mark.vcr()
async def test_rebuild_full_with_accessible_source(temp_db_path):
"""FULL rebuild re-ingests from source when the URI is accessible.
Covers the main path in _rebuild_full (source-accessible branch): the
document is deleted and re-created from its URI, producing a new ID.
"""
async with HaikuRAG(temp_db_path, create=True) as client:
with tempfile.TemporaryDirectory() as temp_dir:
source_path = Path(temp_dir) / "source.txt"
source_path.write_text("Fresh content from an accessible file source.")
original = await client.create_document_from_source(source=source_path)
assert not isinstance(original, list)
assert original.id is not None
original_id = original.id
processed_ids = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)
]
# Original doc was deleted and a new one created; the old ID
# must not appear, and exactly one new ID must have been yielded.
assert original_id not in processed_ids
assert len(processed_ids) == 1
new_doc = await client.get_document_by_id(processed_ids[0])
assert new_doc is not None
assert new_doc.uri == source_path.as_uri()
assert "Fresh content" in new_doc.content
async def test_rebuild_title_only_handles_llm_failure(temp_db_path, monkeypatch):
"""TITLE_ONLY: a failure on one document does not abort the generator.
The first document raises during title generation (simulated LLM error);
the second succeeds. Rebuild must log-and-skip the failure, yield only
the successful document, and persist its new title.
"""
from haiku.rag.store.models.document import Document
async with HaikuRAG(temp_db_path, create=True) as client:
# Skip embedding — TITLE_ONLY only touches documents.
doc1 = await client.document_repository.create(
Document(content="doc one body", metadata={})
)
doc2 = await client.document_repository.create(
Document(content="doc two body", metadata={})
)
assert doc1.id is not None and doc2.id is not None
async def fake_generate_title(doc):
if doc.id == doc1.id:
raise RuntimeError("simulated LLM failure")
return "Second Title"
monkeypatch.setattr(client, "generate_title", fake_generate_title)
processed_ids = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY)
]
assert processed_ids == [doc2.id]
refreshed = await client.get_document_by_id(doc2.id)
assert refreshed is not None
assert refreshed.title == "Second Title"
untouched = await client.get_document_by_id(doc1.id)
assert untouched is not None
assert untouched.title is None
@pytest.mark.vcr()
async def test_rebuild_batch_size_flush(temp_db_path, monkeypatch):
"""RECHUNK flushes in batches and yields every document.
Forces a tiny batch size so three docs trigger at least one mid-loop
flush plus the final flush. Regression guard for the batched-write path
in _rebuild_rechunk.
"""
from haiku.rag.client import rebuild as rebuild_module
monkeypatch.setattr(rebuild_module, "_REBUILD_BATCH_SIZE", 2)
async with HaikuRAG(temp_db_path, create=True) as client:
ids: list[str] = []
for i in range(3):
doc = await client.create_document(content=f"batch flush doc {i}")
assert doc.id is not None
ids.append(doc.id)
processed = [
doc_id async for doc_id in client.rebuild_database(mode=RebuildMode.RECHUNK)
]
assert sorted(processed) == sorted(ids)
for doc_id in ids:
chunks = await client.chunk_repository.get_by_document_id(doc_id)
assert len(chunks) > 0