learningcircuit--local-deep-research/tests/research_library/services/test_library_rag_service_rollback.py

"""Regression tests for session-rollback hygiene in LibraryRAGService.

Follow-up to #3827: the service's `index_document` and
`remove_document_from_rag` methods both open a thread-local session,
perform writes (`session.flush()` / `session.commit()`), and catch
exceptions in an outer ``except`` block. Without an explicit rollback
that catch leaves the shared session in ``PendingRollbackError`` state —
every subsequent ORM operation on the same thread cascades.

These tests force an exception inside the try-block and assert that
``session.rollback`` was called before the helper returned its error
dict.
"""

from unittest.mock import MagicMock, patch


_MOD = "local_deep_research.research_library.services.library_rag_service"


def _make_service(**overrides):
    """Instantiate LibraryRAGService with all heavy deps mocked out."""
    with (
        patch(f"{_MOD}.LocalEmbeddingManager") as _lem,
        patch(f"{_MOD}.get_user_db_session"),
        patch(f"{_MOD}.FileIntegrityManager"),
        patch(f"{_MOD}.get_text_splitter"),
    ):
        _lem.return_value.embeddings = MagicMock()
        from local_deep_research.research_library.services.library_rag_service import (
            LibraryRAGService,
        )

        defaults = dict(username="testuser", db_password="pw")
        defaults.update(overrides)
        return LibraryRAGService(**defaults)


def _ctx_yielding(session):
    """Context-manager mock that yields *session*."""
    ctx = MagicMock()
    ctx.__enter__ = MagicMock(return_value=session)
    ctx.__exit__ = MagicMock(return_value=False)
    return ctx


def _make_doc(text="Long text content for indexing test."):
    doc = MagicMock()
    doc.id = "doc-1"
    doc.text_content = text
    doc.title = "Test"
    doc.filename = None
    doc.original_url = "http://example.com/d"
    doc.authors = None
    doc.published_date = None
    doc.doi = None
    doc.arxiv_id = None
    doc.pmid = None
    doc.pmcid = None
    doc.extraction_method = None
    doc.word_count = None
    return doc


# ---------------------------------------------------------------------------
# index_document: rollback fires when the inner work raises.
# ---------------------------------------------------------------------------


@patch(f"{_MOD}.ensure_in_collection")
@patch(f"{_MOD}.get_user_db_session")
def test_index_document_rolls_back_session_on_exception(
    mock_get_session, mock_ensure
):
    """If text-splitting or any inner step raises, the outer except must
    call session.rollback() before returning the error dict."""
    svc = _make_service()

    session = MagicMock()
    mock_get_session.return_value = _ctx_yielding(session)

    session.query.return_value.filter_by.return_value.first.return_value = (
        _make_doc()
    )
    mock_ensure.return_value = MagicMock(indexed=False, chunk_count=0)

    # Force the inner work to raise — text_splitter is patched in
    # _make_service, but make it explicit here.
    svc.text_splitter = MagicMock()
    svc.text_splitter.split_documents.side_effect = RuntimeError(
        "simulated splitter failure"
    )

    result = svc.index_document("doc-1", "coll-1")

    assert result["status"] == "error"
    session.rollback.assert_called_once()


# ---------------------------------------------------------------------------
# remove_document_from_rag: rollback fires when commit (or anything before
# it) raises.
# ---------------------------------------------------------------------------


@patch(f"{_MOD}.get_user_db_session")
def test_remove_document_from_rag_rolls_back_session_on_exception(
    mock_get_session,
):
    """If session.commit() (or earlier work) raises, the outer except
    must call session.rollback()."""
    svc = _make_service()

    session = MagicMock()
    mock_get_session.return_value = _ctx_yielding(session)

    # First filter_by chain returns a non-None DocumentCollection so we
    # enter the try block.
    doc_collection = MagicMock(indexed=True, chunk_count=3)
    session.query.return_value.filter_by.return_value.first.return_value = (
        doc_collection
    )

    # Make the embedding-manager's chunk deletion raise inside the try
    # block (mirrors a real DB / FAISS layer failure).
    svc.embedding_manager = MagicMock()
    svc.embedding_manager._delete_chunks_from_db.side_effect = RuntimeError(
        "simulated chunk-delete failure"
    )

    result = svc.remove_document_from_rag("doc-1", "coll-1")

    assert result["status"] == "error"
    session.rollback.assert_called_once()