haiku.rag/tests/test_versioning.py

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 in-flight background vacuum tasks to complete."""
    await client._await_vacuum_tasks()


@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_aexit_awaits_all_background_vacuums(temp_db_path, monkeypatch):
    """Multiple create_document calls schedule multiple vacuum tasks; __aexit__
    must await all of them, not just the last-scheduled one.

    Scenario: Task A acquires the vacuum lock and is slow. Task B is scheduled
    while Task A still holds the lock — Task B sees the lock held and returns
    immediately. If the client only tracks the most recently scheduled task,
    __aexit__ awaits the fast no-op B and closes the connection while Task A
    is still running.
    """
    from haiku.rag.config import Config

    monkeypatch.setattr(Config.storage, "auto_vacuum", True)

    first_vacuum_completed = asyncio.Event()

    async with HaikuRAG(db_path=temp_db_path, create=True) as client:
        call_count = 0

        async def slow_vacuum(*_args, **_kwargs):
            nonlocal call_count
            call_count += 1
            my_num = call_count
            # Mimic the real vacuum's skip-if-running behavior.
            if client.store._vacuum_lock.locked():
                return
            async with client.store._vacuum_lock:
                if my_num == 1:
                    # Hold the lock longer than any other operation in the
                    # test so Task A cannot finish incidentally. __aexit__
                    # must explicitly wait for this task.
                    await asyncio.sleep(2.0)
                    first_vacuum_completed.set()

        client.store.vacuum = slow_vacuum

        await client.create_document(content="triggers first vacuum")
        # Let Task A start and acquire the vacuum lock before scheduling B.
        await asyncio.sleep(0.02)
        await client.create_document(content="triggers second vacuum")

    assert first_vacuum_completed.is_set(), (
        "__aexit__ returned before the first vacuum task 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}"
        )


async def test_close_suppresses_failing_drain_vacuum(temp_db_path, monkeypatch):
    """A failing final vacuum on close must not raise out of __aexit__,
    where it would mask an in-flight exception from the context body."""
    client = HaikuRAG(db_path=temp_db_path, create=True)
    await client.__aenter__()

    calls: list[int] = []

    async def boom(*args, **kwargs):
        calls.append(1)
        raise RuntimeError("vacuum boom")

    # A finished task in the set forces the drain branch to run.
    task = asyncio.create_task(asyncio.sleep(0))
    await task
    client._vacuum_tasks.add(task)
    monkeypatch.setattr(client.store, "vacuum", boom)

    # Must not raise despite the drain vacuum erroring.
    await client.__aexit__(None, None, None)

    assert calls, "drain vacuum should have been attempted"