Improve test coverage
This commit is contained in:
Yiorgis Gozadinos
parent
d99e48a9a3
commit
258ffa41fa
5 changed files with 467 additions and 4 deletions
|
|
@ -16,6 +16,8 @@ if TYPE_CHECKING:
|
|||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_REBUILD_BATCH_SIZE = 50
|
||||
|
||||
|
||||
async def rebuild_database(
|
||||
client: "HaikuRAG", mode: "RebuildMode | None" = None
|
||||
|
|
@ -204,7 +206,6 @@ async def _rebuild_rechunk(
|
|||
"""Re-chunk and re-embed from existing document content."""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
batch_size = 50
|
||||
pending_chunks: list[Chunk] = []
|
||||
pending_docs: list[Document] = []
|
||||
pending_doc_ids: list[str] = []
|
||||
|
|
@ -234,7 +235,7 @@ async def _rebuild_rechunk(
|
|||
pending_doc_ids.append(doc.id)
|
||||
|
||||
# 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)
|
||||
for doc_id in pending_doc_ids:
|
||||
yield doc_id
|
||||
|
|
@ -255,7 +256,6 @@ async def _rebuild_full(
|
|||
"""Full rebuild: re-convert from source, re-chunk, re-embed."""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
batch_size = 50
|
||||
pending_chunks: list[Chunk] = []
|
||||
pending_docs: list[Document] = []
|
||||
pending_doc_ids: list[str] = []
|
||||
|
|
@ -312,7 +312,7 @@ async def _rebuild_full(
|
|||
pending_doc_ids.append(doc.id)
|
||||
|
||||
# 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)
|
||||
for doc_id in pending_doc_ids:
|
||||
yield doc_id
|
||||
|
|
|
|||
242
tests/cassettes/test_rebuild/test_rebuild_batch_size_flush.yaml
Normal file
242
tests/cassettes/test_rebuild/test_rebuild_batch_size_flush.yaml
Normal file
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -1272,6 +1272,39 @@ async def test_client_convert_file_not_found(temp_db_path):
|
|||
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()
|
||||
async def test_client_convert_unsupported_extension(temp_db_path):
|
||||
"""Test convert() raises ValueError for unsupported file extension."""
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import TypedDict
|
||||
|
||||
import pytest
|
||||
|
|
@ -263,3 +265,107 @@ async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path):
|
|||
|
||||
# Chunk IDs should change (chunks are recreated)
|
||||
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
|
||||
|
|
|
|||
Loading…
Reference in a new issue