library-desk/tests/test_search_persistence.py

"""
Offline tests for the Phase 6 persistence rewrite (single atomic write,
off the hot path) and its SHAPE CONTRACT with the consolidation service.

The consolidation repair loop (consolidation_service.py) consumes the
persisted graph:

- _find_unprocessed_searches:
    MATCH (sq:SearchQuery {processed: false}) WHERE sq.timestamp > ...
    RETURN sq.id, sq.query, sq.user, sq.timestamp, sq.total_results,
           sq.web_count, sq.keywords
- _get_web_results:
    MATCH (sq:SearchQuery {id: $search_id})-[f:FOUND]->(wr:WebResult)
    RETURN wr.url, wr.title, wr.content, f.rank, f.rrf_score
- _mark_search_processed:
    SET sq.processed = true

These tests pin that the new UNWIND-based persistence still writes every
node property, label, and relationship property that consolidation reads.
"""

import asyncio
import json
import uuid
from unittest.mock import AsyncMock, MagicMock

import pytest

from src.models.hybrid_rag import HybridRAGConfig
from src.services.hybrid_rag_service import HybridRAGService

TENANT = "llm_tester"
TENANT_BASE_LABEL = "User_Llm_Tester"
TENANT_DOC_LABEL = "User_Llm_Tester_Document"


@pytest.fixture
def mock_neo4j():
    neo4j = MagicMock()
    neo4j.execute_query = AsyncMock(return_value=[])
    neo4j.execute_write = AsyncMock(return_value=[{"id": "sid"}])
    return neo4j


@pytest.fixture
def service(mock_neo4j):
    vector = MagicMock()
    response = MagicMock()
    response.results = []
    vector.search = AsyncMock(return_value=response)

    graph = MagicMock()
    graph.neo4j = mock_neo4j
    graph.search_documents = AsyncMock(return_value=[])
    graph.get_related_documents = AsyncMock(return_value=[])

    ollama = MagicMock()
    ollama.generate_text = AsyncMock(
        return_value='{"core_keywords": ["q"], "synonyms": {}}'
    )

    settings = MagicMock()
    settings.ollama_llm_model = "test-model"
    settings.vector_similarity_threshold = 0.7

    return HybridRAGService(
        vector_service=vector,
        graph_service=graph,
        searxng_client=MagicMock(),
        ollama_client=ollama,
        content_extractor=MagicMock(),
        settings=settings,
    )


FINAL_RESULTS = [
    {
        "result": {"page_id": 42, "title": "wiki hit"},
        "source_type": "wiki", "sources": ["vector"],
        "rrf_score": 0.5, "final_rank": 1,
    },
    {
        "result": {
            "url": "http://example.com",
            "title": "web hit",
            "content": "x" * 5000,  # must be truncated to 1000
        },
        "source_type": "web", "sources": ["web"],
        "rrf_score": 0.4, "final_rank": 2,
    },
]


async def _persist(service, **overrides):
    kwargs = dict(
        search_id="sid-123",
        query="test query",
        user=TENANT,
        keywords_data={"core_keywords": ["docker"], "synonyms": {"k8s": ["kubernetes"]}},
        raw_results={"vector": [1], "graph": [], "web": [1, 2]},
        final_results=FINAL_RESULTS,
        timing={"total_ms": 123.0},
    )
    kwargs.update(overrides)
    return await service._persist_search_for_librarian(**kwargs)


@pytest.mark.unit
class TestPersistenceIsAtomic:
    async def test_single_write_transaction_no_autocommit_queries(
        self, service, mock_neo4j
    ):
        result = await _persist(service)

        assert result == "sid-123"
        mock_neo4j.execute_write.assert_awaited_once()
        mock_neo4j.execute_query.assert_not_awaited()

    async def test_failure_returns_none(self, service, mock_neo4j):
        mock_neo4j.execute_write = AsyncMock(side_effect=RuntimeError("boom"))

        assert await _persist(service) is None


@pytest.mark.unit
class TestConsolidationShapeContract:
    """Every property/label/relationship consolidation reads must be written."""

    async def test_searchquery_node_shape(self, service, mock_neo4j):
        await _persist(service)
        cypher, params = mock_neo4j.execute_write.await_args.args[:2]

        # Interoperable label + tenant label (consolidation matches bare
        # :SearchQuery, tenant scoping needs the prefixed label)
        assert f":{TENANT_BASE_LABEL}_SearchQuery:SearchQuery" in cypher
        # _find_unprocessed_searches filters on these
        assert "processed: false" in cypher
        assert "timestamp: datetime()" in cypher
        # ... and returns these properties
        for prop in ("id", "query", "user", "total_results", "web_count", "keywords"):
            assert f"{prop}: ${'search_id' if prop == 'id' else prop}" in cypher, prop
        assert params["search_id"] == "sid-123"
        assert params["query"] == "test query"
        assert params["user"] == TENANT
        assert params["total_results"] == 2
        assert params["web_count"] == 2
        assert params["keywords"] == ["docker"]
        assert json.loads(params["synonyms"]) == {"k8s": ["kubernetes"]}

    async def test_webresult_shape_and_found_relationship(self, service, mock_neo4j):
        await _persist(service)
        cypher, params = mock_neo4j.execute_write.await_args.args[:2]

        # _get_web_results traverses (sq)-[f:FOUND]->(wr:WebResult) and
        # reads wr.url, wr.title, wr.content, f.rank, f.rrf_score
        assert f":{TENANT_BASE_LABEL}_WebResult:WebResult" in cypher
        for fragment in ("url: wl.url", "title: wl.title", "content: wl.content"):
            assert fragment in cypher, fragment
        assert "rank: wl.rank" in cypher
        assert "rrf_score: wl.rrf_score" in cypher
        assert 'source: "web"' in cypher

        web = params["web_links"]
        assert len(web) == 1
        assert web[0]["url"] == "http://example.com"
        assert web[0]["rank"] == 1
        assert web[0]["rrf_score"] == 0.4
        assert len(web[0]["content"]) == 1000  # truncation preserved

    async def test_document_links_tenant_scoped(self, service, mock_neo4j):
        await _persist(service)
        cypher, params = mock_neo4j.execute_write.await_args.args[:2]

        assert f"MATCH (d:{TENANT_DOC_LABEL}:Document {{page_id: link.page_id}})" in cypher
        assert "MERGE (sq)-[f:FOUND]->(d)" in cypher
        for fragment in ("f.source = link.source", "f.rank = link.rank",
                         "f.rrf_score = link.rrf_score",
                         "f.final_rank = link.final_rank"):
            assert fragment in cypher, fragment

        docs = params["doc_links"]
        assert docs == [{
            "page_id": 42, "source": "wiki", "rank": 1,
            "rrf_score": 0.5, "final_rank": 1,
        }]

    async def test_empty_doc_links_cannot_swallow_web_results(
        self, service, mock_neo4j
    ):
        """UNWIND [] yields no rows; the CALL subqueries must isolate that."""
        await _persist(service, final_results=[FINAL_RESULTS[1]])
        cypher, params = mock_neo4j.execute_write.await_args.args[:2]

        assert params["doc_links"] == []
        assert len(params["web_links"]) == 1
        # Both UNWINDs live in aggregating CALL subqueries
        assert cypher.count("CALL {") == 2


@pytest.mark.unit
class TestPersistenceOffHotPath:
    async def test_search_returns_upfront_id_and_persists_in_background(
        self, service, mock_neo4j
    ):
        config = HybridRAGConfig(
            enable_vector=True, enable_graph=False, enable_web=False,
            enable_volatile=False, enable_documents=False,
            enable_reranking=False, enable_enrichment=False,
        )

        response = await service.search("q", TENANT, config)

        # search_id is generated up front and returned immediately
        assert response.search_id
        uuid.UUID(response.search_id)  # valid uuid4
        assert response.timing.persistence_ms == 0.0

        # The write happens in a background task, not on the request path
        pending = list(service._background_tasks)
        assert len(pending) == 1
        await asyncio.gather(*pending)

        mock_neo4j.execute_write.assert_awaited_once()
        params = mock_neo4j.execute_write.await_args.args[1]
        assert params["search_id"] == response.search_id