276 lines
11 KiB
Python
276 lines
11 KiB
Python
import asyncio
|
|
|
|
import pytest
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.store.engine import Store
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_version_rollback_on_create_failure(temp_db_path):
|
|
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
|
# Patch chunk_repository.create to succeed then fail, triggering rollback
|
|
orig_create = client.chunk_repository.create
|
|
|
|
async def succeed_then_fail(chunks):
|
|
await orig_create(chunks)
|
|
raise RuntimeError("boom")
|
|
|
|
client.chunk_repository.create = succeed_then_fail
|
|
|
|
# Attempt to create document; expect failure and rollback
|
|
content = "Hello, rollback!"
|
|
|
|
with pytest.raises(RuntimeError):
|
|
await client.create_document(content=content)
|
|
|
|
# State should be restored (no documents/chunks)
|
|
docs = await client.list_documents()
|
|
assert len(docs) == 0
|
|
all_chunks = await client.chunk_repository.list_all()
|
|
assert len(all_chunks) == 0
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_version_rollback_on_update_failure(temp_db_path):
|
|
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
|
# Create a valid document first
|
|
base_content = "Base content"
|
|
created = await client.create_document(content=base_content)
|
|
|
|
# Patch chunk_repository.create to succeed then fail during update
|
|
orig_create = client.chunk_repository.create
|
|
|
|
async def succeed_then_fail(chunks):
|
|
await orig_create(chunks)
|
|
raise RuntimeError("update fail")
|
|
|
|
client.chunk_repository.create = succeed_then_fail
|
|
|
|
# Attempt update
|
|
with pytest.raises(RuntimeError):
|
|
await client.update_document(
|
|
document_id=created.id,
|
|
content="Updated content",
|
|
)
|
|
|
|
# Content and chunks should remain the original
|
|
persisted = await client.get_document_by_id(created.id)
|
|
assert persisted is not None
|
|
assert persisted.content == base_content
|
|
original_chunks = await client.chunk_repository.get_by_document_id(created.id)
|
|
assert len(original_chunks) > 0
|
|
|
|
|
|
async def test_new_database_does_not_run_upgrades(monkeypatch, temp_db_path):
|
|
def fail_if_called(*_args, **_kwargs):
|
|
raise AssertionError("run_pending_upgrades should not be called for new DB")
|
|
|
|
monkeypatch.setattr(
|
|
"haiku.rag.store.upgrades.run_pending_upgrades",
|
|
fail_if_called,
|
|
)
|
|
|
|
async with Store(temp_db_path, create=True):
|
|
pass
|
|
|
|
|
|
async def test_existing_database_checks_migrations(monkeypatch, temp_db_path):
|
|
async with Store(temp_db_path, create=True):
|
|
pass
|
|
|
|
from haiku.rag.store import upgrades
|
|
|
|
called = {"value": False}
|
|
original_get_pending = upgrades.get_pending_upgrades
|
|
|
|
def mark_called(*args, **kwargs):
|
|
called["value"] = True
|
|
return original_get_pending(*args, **kwargs)
|
|
|
|
monkeypatch.setattr(
|
|
"haiku.rag.store.upgrades.get_pending_upgrades",
|
|
mark_called,
|
|
)
|
|
|
|
# Opening an existing database should check for pending migrations
|
|
async with Store(temp_db_path):
|
|
pass
|
|
|
|
assert called["value"]
|
|
|
|
|
|
async def _wait_for_background_vacuum(client):
|
|
"""Wait for any background vacuum task to complete."""
|
|
if client._vacuum_task is not None and not client._vacuum_task.done():
|
|
await client._vacuum_task
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_vacuum_with_retention_threshold(temp_db_path):
|
|
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
|
# Create first document
|
|
await client.create_document(content="First document")
|
|
await _wait_for_background_vacuum(client)
|
|
|
|
# Create second document
|
|
await client.create_document(content="Second document")
|
|
await _wait_for_background_vacuum(client)
|
|
|
|
store = client.store
|
|
|
|
# Get initial version counts (should have multiple versions from creates)
|
|
initial_doc_versions = len(await store.documents_table.list_versions())
|
|
initial_chunk_versions = len(await store.chunks_table.list_versions())
|
|
|
|
assert initial_doc_versions > 1, "Should have multiple document table versions"
|
|
assert initial_chunk_versions > 1, "Should have multiple chunk table versions"
|
|
|
|
# Vacuum with default threshold (60 seconds) - should keep recent versions
|
|
# Note: vacuum may create new versions even when not cleaning up old ones
|
|
await store.vacuum()
|
|
|
|
after_default_doc_versions = len(await store.documents_table.list_versions())
|
|
after_default_chunk_versions = len(await store.chunks_table.list_versions())
|
|
|
|
# After vacuum with retention, version count should stay the same or increase
|
|
# (optimize may create new versions) but not decrease
|
|
assert after_default_doc_versions >= initial_doc_versions, (
|
|
"Default vacuum should not remove recent versions"
|
|
)
|
|
assert after_default_chunk_versions >= initial_chunk_versions, (
|
|
"Default vacuum should not remove recent versions"
|
|
)
|
|
|
|
# Vacuum with 0 threshold - should significantly reduce versions
|
|
await store.vacuum(retention_seconds=0)
|
|
|
|
after_zero_doc_versions = len(await store.documents_table.list_versions())
|
|
after_zero_chunk_versions = len(await store.chunks_table.list_versions())
|
|
|
|
# After aggressive vacuum, should have minimal versions (1-2)
|
|
# Note: optimize operation may create a version after cleanup
|
|
assert after_zero_doc_versions <= 2, (
|
|
f"Should have minimal document versions after vacuum(0), got {after_zero_doc_versions}"
|
|
)
|
|
assert after_zero_chunk_versions <= 2, (
|
|
f"Should have minimal chunk versions after vacuum(0), got {after_zero_chunk_versions}"
|
|
)
|
|
|
|
# And it should be significantly fewer than before
|
|
assert after_zero_doc_versions < initial_doc_versions, (
|
|
"Should have fewer versions after vacuum(0)"
|
|
)
|
|
assert after_zero_chunk_versions < initial_chunk_versions, (
|
|
"Should have fewer versions after vacuum(0)"
|
|
)
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_vacuum_completes_before_context_exit(temp_db_path, monkeypatch):
|
|
"""Test that background vacuum completes when context manager exits."""
|
|
from haiku.rag.config import Config
|
|
|
|
# Set aggressive vacuum retention for this test
|
|
monkeypatch.setattr(Config.storage, "vacuum_retention_seconds", 0)
|
|
|
|
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
|
# Create multiple documents - each creation triggers automatic vacuum with retention=0
|
|
# This aggressively cleans up old versions between operations
|
|
for i in range(3):
|
|
await client.create_document(content=f"Test document {i}")
|
|
|
|
# After context exit, automatic vacuum should have kept versions minimal
|
|
async with Store(temp_db_path, create=True) as store:
|
|
final_versions = len(await store.documents_table.list_versions())
|
|
|
|
# With retention_seconds=0, vacuum aggressively cleans up between operations
|
|
# Should have very few versions remaining (1-2)
|
|
assert final_versions <= 2, (
|
|
f"Aggressive vacuum should keep minimal versions, got {final_versions}"
|
|
)
|
|
assert final_versions >= 1, "Should have at least one version remaining"
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_aexit_awaits_background_vacuum(temp_db_path, monkeypatch):
|
|
"""__aexit__ must await any in-flight background vacuum, not just release the lock.
|
|
|
|
Background vacuum runs as an asyncio task; the event loop may not have scheduled
|
|
it yet when __aexit__ runs. Simply acquiring the vacuum lock (which is free until
|
|
the task actually starts) would let close() proceed before vacuum runs.
|
|
"""
|
|
from haiku.rag.config import Config
|
|
|
|
monkeypatch.setattr(Config.storage, "auto_vacuum", True)
|
|
|
|
vacuum_started = asyncio.Event()
|
|
vacuum_completed = asyncio.Event()
|
|
|
|
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
|
original_vacuum = client.store.vacuum
|
|
|
|
async def instrumented_vacuum(*args, **kwargs):
|
|
vacuum_started.set()
|
|
# Delay so __aexit__ would see an unstarted/incomplete task if it
|
|
# relied on the lock rather than awaiting the task directly.
|
|
await asyncio.sleep(0.05)
|
|
await original_vacuum(*args, **kwargs)
|
|
vacuum_completed.set()
|
|
|
|
client.store.vacuum = instrumented_vacuum
|
|
|
|
await client.create_document(content="triggers background vacuum")
|
|
|
|
assert vacuum_started.is_set(), "Background vacuum task never ran"
|
|
assert vacuum_completed.is_set(), "__aexit__ exited before vacuum finished"
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_auto_vacuum_disabled_skips_vacuum(temp_db_path, monkeypatch):
|
|
"""Test that auto_vacuum=False prevents automatic vacuum after operations."""
|
|
from haiku.rag.config import Config
|
|
|
|
# Disable auto-vacuum
|
|
monkeypatch.setattr(Config.storage, "auto_vacuum", False)
|
|
|
|
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
|
# Create multiple documents
|
|
for i in range(3):
|
|
await client.create_document(content=f"Test document {i}")
|
|
|
|
# Count versions - should accumulate without vacuum
|
|
doc_versions = len(await client.store.documents_table.list_versions())
|
|
chunk_versions = len(await client.store.chunks_table.list_versions())
|
|
|
|
# Without auto-vacuum, versions should accumulate (more than 3 from creates)
|
|
assert doc_versions >= 3, (
|
|
f"Without auto-vacuum, should have accumulated versions, got {doc_versions}"
|
|
)
|
|
assert chunk_versions >= 3, (
|
|
f"Without auto-vacuum, should have accumulated versions, got {chunk_versions}"
|
|
)
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_auto_vacuum_enabled_triggers_vacuum(temp_db_path, monkeypatch):
|
|
"""Test that auto_vacuum=True (default) triggers vacuum after operations."""
|
|
from haiku.rag.config import Config
|
|
|
|
# Enable auto-vacuum with aggressive retention
|
|
monkeypatch.setattr(Config.storage, "auto_vacuum", True)
|
|
monkeypatch.setattr(Config.storage, "vacuum_retention_seconds", 0)
|
|
|
|
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
|
# Create multiple documents
|
|
for i in range(3):
|
|
await client.create_document(content=f"Test document {i}")
|
|
|
|
# After context exit, vacuum should have cleaned up
|
|
async with Store(temp_db_path, create=True) as store:
|
|
final_versions = len(await store.documents_table.list_versions())
|
|
|
|
# With auto_vacuum=True and retention=0, should have minimal versions
|
|
assert final_versions <= 2, (
|
|
f"With auto-vacuum enabled, should have minimal versions, got {final_versions}"
|
|
)
|